Update main.py

main.py

@@ -1,146 +1,125 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
+import os
+import tempfile
+from fastapi import FastAPI, HTTPException, BackgroundTasks
 from pydantic import BaseModel
-from fastapi.middleware.cors import CORSMiddleware
-import uuid
-from concurrent.futures import ThreadPoolExecutor
-from pymongo import MongoClient
-from urllib.parse import quote_plus
-from langchain_groq import ChatGroq
-from aura_sr import AuraSR
-from io import BytesIO
 from PIL import Image
+import io
 import requests
-import os
+from typing import Optional
+import threading
+import uuid
 
 app = FastAPI()
 
-# Middleware for CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Globals
-executor = ThreadPoolExecutor(max_workers=10)
-llm = None
-upscale_model = None
-client = MongoClient(f"mongodb+srv://hammad:{quote_plus('momimaad@123')}@cluster0.2a9yu.mongodb.net/")
-db = client["Flux"]
-collection = db["chat_histories"]
-chat_sessions = {}
-image_storage_dir = "./images"  # Directory to save images locally
-
-# Ensure the image storage directory exists
-os.makedirs(image_storage_dir, exist_ok=True)
-
-@app.on_event("startup")
-async def startup():
-    global llm, upscale_model
-    llm = ChatGroq(
-        model="llama-3.3-70b-versatile",
-        temperature=0.7,
-        max_tokens=1024,
-        api_key="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",
-    )
-    upscale_model = AuraSR.from_pretrained("fal/AuraSR-v2")
+# Use a temporary directory for storing images
+temp_dir = tempfile.mkdtemp()
+print(f"Temporary directory for images: {temp_dir}")
+
+# Ensure that the temporary directory exists
+os.makedirs(temp_dir, exist_ok=True)
+
+
+# Endpoint models
+class NewChatRequest(BaseModel):
+    user_name: str
+
+
+class GenerateImageRequest(BaseModel):
+    prompt: str
+    width: int = 1024
+    height: int = 1024
+    inference_steps: int = 50
+    guidance_scale: Optional[float] = 7.5
 
-@app.on_event("shutdown")
-def shutdown():
-    client.close()
-    executor.shutdown()
 
-# Pydantic models
-class ImageRequest(BaseModel):
+class UpscaleImageRequest(BaseModel):
+    image_id: str
+
+
+class SetPromptRequest(BaseModel):
     subject: str
     style: str
     color_theme: str
     elements: str
-    color_mode: str
     lighting_conditions: str
     framing_style: str
     material_details: str
     text: str
     background_details: str
-    user_prompt: str
     chat_id: str
+    user_prompt: str
+
 
-# Helper Functions
-def generate_chat_id():
+# Dictionary to store images temporarily
+image_store = {}
+
+
+@app.post("/new-chat")
+def new_chat(request: NewChatRequest):
     chat_id = str(uuid.uuid4())
-    chat_sessions[chat_id] = collection
-    return chat_id
-
-def get_chat_history(chat_id):
-    messages = collection.find({"session_id": chat_id})
-    return "\n".join(
-        f"User: {msg['content']}" if msg['role'] == "user" else f"AI: {msg['content']}"
-        for msg in messages
-    )
-
-def save_image_locally(image, filename):
-    filepath = os.path.join(image_storage_dir, filename)
-    image.save(filepath, format="PNG")
-    return filepath
-
-# Endpoints
-@app.post("/new-chat", response_model=dict)
-async def new_chat():
-    chat_id = generate_chat_id()
-    return {"chat_id": chat_id}
-
-@app.post("/generate-image", response_model=dict)
-async def generate_image(request: ImageRequest, background_tasks: BackgroundTasks):
-    def process_request():
-        chat_history = get_chat_history(request.chat_id)
-        prompt = f"""
-        Subject: {request.subject}
-        Style: {request.style}
-        ...
-        Chat History: {chat_history}
-        User Prompt: {request.user_prompt}
-        """
-        refined_prompt = llm.invoke(prompt).content.strip()
-        collection.insert_one({"session_id": request.chat_id, "role": "user", "content": request.user_prompt})
-        collection.insert_one({"session_id": request.chat_id, "role": "ai", "content": refined_prompt})
+    # Simulate creating a chat session
+    return {"chat_id": chat_id, "message": f"New chat created for user: {request.user_name}"}
 
+
+@app.post("/generate-image")
+def generate_image(request: GenerateImageRequest, background_tasks: BackgroundTasks):
+    try:
         # Simulate image generation
-        response = requests.post(
-            "https://api.bfl.ml/v1/flux-pro-1.1",
-            json={"prompt": refined_prompt}
-        ).json()
-        image_url = response["result"]["sample"]
-
-        # Download and save the image locally
-        image_response = requests.get(image_url)
-        img = Image.open(BytesIO(image_response.content))
-        filename = f"generated_{uuid.uuid4()}.png"
-        filepath = save_image_locally(img, filename)
-        return filepath
-
-    task = executor.submit(process_request)
-    background_tasks.add_task(task)
-    return {"status": "Processing"}
-
-@app.post("/upscale-image", response_model=dict)
-async def upscale_image(image_url: str, background_tasks: BackgroundTasks):
-    def process_image():
-        response = requests.get(image_url)
-        img = Image.open(BytesIO(response.content))
-        upscaled_image = upscale_model.upscale_4x_overlapped(img)
-        filename = f"upscaled_{uuid.uuid4()}.png"
-        filepath = save_image_locally(upscaled_image, filename)
-        return filepath
-
-    task = executor.submit(process_image)
-    background_tasks.add_task(task)
-    return {"status": "Processing"}
-
-@app.post("/set-prompt", response_model=dict)
-async def set_prompt(chat_id: str, user_prompt: str):
-    chat_history = get_chat_history(chat_id)
-    refined_prompt = llm.invoke(f"{chat_history}\nUser Prompt: {user_prompt}").content.strip()
-    collection.insert_one({"session_id": chat_id, "role": "user", "content": user_prompt})
-    collection.insert_one({"session_id": chat_id, "role": "ai", "content": refined_prompt})
-    return {"refined_prompt": refined_prompt}
+        image_id = str(uuid.uuid4())
+        image_path = os.path.join(temp_dir, f"{image_id}.png")
+
+        # Mock image generation (replace with actual logic)
+        img = Image.new("RGB", (request.width, request.height), color=(255, 255, 255))
+        img.save(image_path)
+
+        # Store image path
+        image_store[image_id] = image_path
+        return {"image_id": image_id, "message": "Image generation started"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/upscale-image")
+def upscale_image(request: UpscaleImageRequest):
+    image_path = image_store.get(request.image_id)
+    if not image_path:
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    try:
+        # Simulate image upscaling
+        img = Image.open(image_path)
+        upscaled_img = img.resize((img.width * 2, img.height * 2), Image.LANCZOS)
+        upscaled_image_path = os.path.join(temp_dir, f"{request.image_id}_upscaled.png")
+        upscaled_img.save(upscaled_image_path)
+
+        # Update store
+        image_store[request.image_id] = upscaled_image_path
+        return {"image_id": request.image_id, "message": "Image upscaled"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/set-prompt")
+def set_prompt(request: SetPromptRequest):
+    try:
+        # Simulate prompt generation (replace with actual LLM call)
+        refined_prompt = (
+            f"Subject: {request.subject}, Style: {request.style}, "
+            f"Color Theme: {request.color_theme}, Elements: {request.elements}, "
+            f"Background: {request.background_details}"
+        )
+        return {"refined_prompt": refined_prompt}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e)}
+
+
+# Clean up the temporary directory on shutdown
+@app.on_event("shutdown")
+def cleanup():
+    print("Cleaning up temporary directory...")
+    try:
+        for file in os.listdir(temp_dir):
+            os.remove(os.path.join(temp_dir, file))
+        os.rmdir(temp_dir)
+    except Exception as e:
+        print(f"Error cleaning up temporary directory: {e}")