Update app.py

app.py

@@ -1,50 +1,44 @@
-python
-import io
-import base64
+from typing import Dict, List
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from diffusers import StableDiffusionXLPipeline
 import torch
-from PIL import Image
-from flask import Flask, request, jsonify
-from diffusers import StableDiffusionPipeline
+import base64
+from io import BytesIO
 
-app = Flask(__name__)
+app = FastAPI()
 
 # Load the model
 model_name = "colt12/maxcushion"
-try:
-    print("Loading model...")
-    pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
-    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-    print("Model loaded successfully.")
-except Exception as e:
-    print(f"Error loading model: {str(e)}")
-    raise
-
-def generate_image(prompt):
-    with torch.no_grad():
-        image = pipe(prompt).images[0]
-    
-    buffered = io.BytesIO()
-    image.save(buffered, format="PNG")
-    image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    
-    return image_base64
+pipe = StableDiffusionXLPipeline.from_pretrained(model_name, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-@app.route('/', methods=['GET'])
-def home():
-    return "Welcome to the Image Generation API. Use the /generate endpoint to generate images from prompts."
+class Item(BaseModel):
+    prompt: str
+    negative_prompt: str = ""
+    num_inference_steps: int = 30
+    guidance_scale: float = 7.5
 
-@app.route('/generate', methods=['POST'])
-def run():
-    if 'prompt' not in request.json:
-        return jsonify({"error": "No prompt provided"}), 400
-    
-    prompt = request.json['prompt']
-    
+@app.post("/generate")
+async def generate(item: Item) -> Dict[str, str]:
     try:
-        result = generate_image(prompt)
-        return jsonify({"image": result})
+        # Generate the image
+        image = pipe(
+            prompt=item.prompt,
+            negative_prompt=item.negative_prompt,
+            num_inference_steps=item.num_inference_steps,
+            guidance_scale=item.guidance_scale
+        ).images[0]
+
+        # Convert to base64
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        return {"image": image_base64}
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        raise HTTPException(status_code=500, detail=str(e))
 
-if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=5000)
+@app.get("/")
+async def root():
+    return {"message": "SDXL Image Generation API"}