Update app.py

app.py

@@ -1,4 +1,6 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Form  # Add Form import
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import logging
 import torch
@@ -16,8 +18,18 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
+
+# Add CORS middleware to allow requests from Framer
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, restrict to your Framer domain
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 pipe = None
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def initialize_pipeline():
     global pipe
@@ -28,11 +40,11 @@ def initialize_pipeline():
         pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
             model_id,
             scheduler=scheduler,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
             low_cpu_mem_usage=True
         )
         pipe = pipe.to(device)
-        logger.info("Stable Diffusion pipeline initialized successfully.")
+        logger.info(f"Stable Diffusion pipeline initialized successfully on {device}.")
     except Exception as e:
         logger.error(f"Failed to initialize pipeline: {str(e)}", exc_info=True)
         raise
@@ -98,7 +110,7 @@ def overlay_face(generated_img, face_img, face_coords):
 def image_to_base64(image: Image.Image) -> str:
     try:
         buffered = BytesIO()
-        image.save(buffered, format="PNG")
+        image.save(buffered, format="JPEG")  # Changed to JPEG to match Framer client expectation
         img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
         logger.info("Image converted to base64 successfully.")
         return img_base64
@@ -111,21 +123,33 @@ async def predict(
     prompt: str = Form(...),
     image: UploadFile = File(...),
     negative_prompt: str = Form("low quality, blurry"),
-    seed: int = Form(66),
-    guidance_scale: float = Form(7.5),
-    num_inference_steps: int = Form(10),
-    strength: float = Form(0.75)
+    seed: str = Form("66"),  # Changed to str to match Framer client
+    guidance_scale: str = Form("7.5"),  # Changed to str to match Framer client
+    num_inference_steps: str = Form("10"),  # Changed to str to match Framer client
+    strength: str = Form("0.75")  # Changed to str to match Framer client
 ):
     global pipe
     try:
         if pipe is None:
+            logger.error("Pipeline not initialized.")
             raise HTTPException(status_code=500, detail="Pipeline not initialized.")
 
         logger.info(f"Received inference request with prompt: {prompt}")
 
+        # Convert string parameters to appropriate types
+        try:
+            seed = int(seed)
+            guidance_scale = float(guidance_scale)
+            num_inference_steps = int(num_inference_steps)
+            strength = float(strength)
+        except ValueError as e:
+            logger.error(f"Invalid parameter format: {str(e)}")
+            raise HTTPException(status_code=400, detail=f"Invalid parameter format: {str(e)}")
+
         # Load and process uploaded image
         logger.info("Loading uploaded image...")
-        ref_image = Image.open(image.file).convert("RGB")
+        image_data = await image.read()
+        ref_image = Image.open(BytesIO(image_data)).convert("RGB")
 
         # Extract face
         logger.info("Extracting face...")
@@ -159,11 +183,9 @@ async def predict(
         result_base64 = image_to_base64(final_img)
 
         logger.info("Inference completed successfully.")
-        return {
-            "original_image": "uploaded_file",
-            "prompt": prompt,
-            "result_image": f"data:image/png;base64,{result_base64}"
-        }
+        return JSONResponse({
+            "result_image": f"data:image/jpeg;base64,{result_base64}"  # Match Framer client expectation
+        })
     except HTTPException as e:
         logger.error(f"HTTP Exception: {str(e)}")
         raise