Update app.py

app.py

@@ -1,26 +1,82 @@
-FROM nvidia/cuda:11.6.2-base-ubuntu20.04
+import torch
+from diffusers import StableDiffusionPipeline
+from fastapi import FastAPI, Response
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+import io
+import base64
+from typing import Optional
+import uvicorn
 
-ENV DEBIAN_FRONTEND=noninteractive
-WORKDIR /code
+# Initialize FastAPI app
+app = FastAPI(title="Stable Diffusion API")
 
-# Install Python and dependencies
-RUN apt-get update && \
-    apt-get install -y python3 python3-pip git && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+# Define input model
+class TextToImageRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    num_inference_steps: Optional[int] = 50
+    guidance_scale: Optional[float] = 7.5
+    height: Optional[int] = 512
+    width: Optional[int] = 512
+    seed: Optional[int] = None
 
-# Copy requirements and install
-COPY requirements.txt .
-RUN pip3 install --no-cache-dir -r requirements.txt
+# Load the model (will be loaded when the Space is initialized)
+model_id = "CompVis/stable-diffusion-v1-4"
 
-# Copy app files
-COPY app.py .
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
+pipe = pipe.to(device)
 
-# Set environment variables for HF
-ENV HF_HOME=/code
+@app.get("/")
+def read_root():
+    return {"message": "Stable Diffusion API is running. Use POST /generate endpoint."}
 
-# Expose port
-EXPOSE 7860
+@app.post("/generate")
+async def generate_image(request: TextToImageRequest):
+    try:
+        # Set seed if provided
+        if request.seed is not None:
+            generator = torch.Generator(device=device).manual_seed(request.seed)
+        else:
+            generator = None
+        
+        # Generate image
+        image = pipe(
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt,
+            num_inference_steps=request.num_inference_steps,
+            guidance_scale=request.guidance_scale,
+            height=request.height,
+            width=request.width,
+            generator=generator
+        ).images[0]
+        
+        # Convert to base64
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        
+        return JSONResponse({
+            "status": "success",
+            "image": img_str,
+            "parameters": {
+                "prompt": request.prompt,
+                "negative_prompt": request.negative_prompt,
+                "steps": request.num_inference_steps,
+                "guidance_scale": request.guidance_scale,
+                "dimensions": f"{request.width}x{request.height}",
+                "seed": request.seed
+            }
+        })
+    
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={"status": "error", "message": str(e)}
+        )
 
-# Run the application with proper host to expose externally
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+# For local testing, not necessary in Spaces
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)