Update main.py

main.py

@@ -1,125 +1,161 @@
-import os
-import tempfile
-from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi import FastAPI, HTTPException, BackgroundTasks, Depends
 from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+import uuid
+import tempfile
+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
-from typing import Optional
-import threading
-import uuid
+import os
 
 app = FastAPI()
 
-# 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)
-
+# 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 = {}
 
-# 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
+# Use a temporary directory for storing images
+image_storage_dir = tempfile.mkdtemp()
+print(f"Temporary directory for images: {image_storage_dir}")
 
+# Dictionary to store images temporarily
+image_store = {}
 
-class UpscaleImageRequest(BaseModel):
-    image_id: str
+@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")
 
+@app.on_event("shutdown")
+def shutdown():
+    client.close()
+    executor.shutdown()
 
-class SetPromptRequest(BaseModel):
+# Pydantic models
+class ImageRequest(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
-    chat_id: str
     user_prompt: str
+    chat_id: str
 
-
-# Dictionary to store images temporarily
-image_store = {}
-
-
-@app.post("/new-chat")
-def new_chat(request: NewChatRequest):
+# Helper Functions
+def generate_chat_id():
     chat_id = str(uuid.uuid4())
-    # Simulate creating a chat session
-    return {"chat_id": chat_id, "message": f"New chat created for user: {request.user_name}"}
-
+    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"""
+        You are a professional assistant responsible for crafting a clear and visually compelling prompt for an image generation model. Your task is to generate a high-quality prompt for creating both the **main subject** and the **background** of the image.
+
+        Image Specifications:
+        - **Subject**: Focus on **{subject}**, highlighting its defining features, expressions, and textures.
+        - **Style**: Emphasize the **{style}**, capturing its key characteristics.
+        - **Background**: Create a background with **{background_details}** that complements and enhances the subject. Ensure it aligns with the color theme and overall composition.
+        - **Camera and Lighting**:
+          - Lighting: Apply **{lighting_conditions}**, emphasizing depth, highlights, and shadows to accentuate the subject and harmonize the background.
+        - **Framing**: Use a **{framing_style}** to enhance the composition around both the subject and the background.
+        - **Materials**: Highlight textures like **{material_details}**, with realistic details and natural imperfections on the subject and background.
+        - **Key Elements**: Include **{elements}** to enrich the subject’s details and add cohesive elements to the background.
+        - **Color Theme**: Follow the **{color_theme}** to set the mood and tone for the entire scene.
+        - Negative Prompt: Avoid grainy, blurry, or deformed outputs.
+        - **Text to Include in Image**: Clearly display the text **"{text}"** as part of the composition (e.g., on a card, badge, or banner) attached to the subject in a realistic and contextually appropriate way.
+        - Write the prompt only in the output. Do not include anything else except the prompt. Do not write Generated Image Prompt: in the output.
+        
+        Chat History:
+        {chat_history_text}
+        
+        User Prompt:
+        {user_prompt}
+        
+        Generated Image 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})
 
-@app.post("/generate-image")
-def generate_image(request: GenerateImageRequest, background_tasks: BackgroundTasks):
-    try:
         # Simulate image generation
-        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}")
+        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"}