Update app.py

app.py

@@ -1,11 +1,9 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import requests
-import io
 import random
 import os
-from PIL import Image
-import base64
+from fastapi.responses import Response
 
 app = FastAPI()
 
@@ -18,7 +16,7 @@ class ImageRequest(BaseModel):
     negative_prompt: str = "(deformed, distorted, disfigured), poorly drawn, bad anatomy"
     steps: int = 4
     cfg_scale: float = 7.0
-    sampler: str = "DPM++ SDE Karras"
+    sampler: str = "DPM++ 2M Karras"
     seed: int = -1
     strength: float = 0.7
 
@@ -44,19 +42,12 @@ def query(prompt: str, negative_prompt: str, steps: int, cfg_scale: float,
     if response.status_code != 200:
         raise HTTPException(status_code=response.status_code, detail=response.text)
     
-    image_bytes = response.content
-    image = Image.open(io.BytesIO(image_bytes))
-    
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    
-    return {"image": f"data:image/png;base64,{img_str}"}
+    return response.content
 
 @app.post("/generate")
 async def generate_image(request: ImageRequest):
     try:
-        result = query(
+        raw_data = query(
             prompt=request.prompt,
             negative_prompt=request.negative_prompt,
             steps=request.steps,
@@ -65,7 +56,7 @@ async def generate_image(request: ImageRequest):
             seed=request.seed,
             strength=request.strength
         )
-        return result
+        return Response(content=raw_data, media_type="application/octet-stream")
     except HTTPException as e:
         raise e
     except Exception as e: