Update app.py

app.py

@@ -1,230 +1,78 @@
-import io
-import os
-import time
-import json
-
-import numpy as np
-from fastapi import FastAPI, HTTPException, Body
-from fastapi.responses import StreamingResponse, Response, HTMLResponse
-from fastapi.middleware import Middleware
-from fastapi.middleware.gzip import GZipMiddleware
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-
-from onnxruntime import InferenceSession
-from huggingface_hub import snapshot_download
-from scipy.io.wavfile import write as write_wav
-
+from io import BytesIO
 from diffusers import OnnxStableDiffusionPipeline
+from huggingface_hub import snapshot_download
 from PIL import Image
+import os
+
+app = FastAPI()
+
+# Global variable to hold the loaded pipeline
+pipeline = None
+model_id = "clip.opt/model.onnx" # Or any other ONNX compatible Stable Diffusion model
+repo_id = "black-forest-labs/FLUX.1-dev-onnx" # Directory to store downloaded ONNX models using snapshot_download
 
 class ImageRequest(BaseModel):
     prompt: str
     num_inference_steps: int = 50
     guidance_scale: float = 7.5
-    format: str = "png" # or "jpeg"
-
-
-model_repo = "runwayml/stable-diffusion-v1-5" # Or any other ONNX compatible Stable Diffusion model
-model_name = "model_q4.onnx" # if specific model file needed, otherwise directory is enough
-voice_file_pattern = "*.bin" # not used, keep for inspiration, remove if not needed
-local_dir = "sd_onnx_models_snapshot" # different folder for sd models
-snapshot_download(
-    repo_id=model_repo,
-    revision="onnx",
-    local_dir=local_dir,
-    local_dir_use_symlinks=False,
-    allow_patterns=["*.onnx", "*.json", "vae/*.onnx"] # Specify necessary file patterns (adjust as needed)
-)
-
-
-pipeline = OnnxStableDiffusionPipeline.from_pretrained(
-    local_dir,  # Use the local path from snapshot_download
-    provider="CPUExecutionProvider", # Or "CUDAExecutionProvider" if you have GPU
-)
-
-
-app = FastAPI(
-    title="FastAPI Image Generation with ONNX",
-    middleware=[Middleware(GZipMiddleware, compresslevel=9)] # maybe compression is not needed for images? check later
-)
-
-
-@app.post("/generate-image/streaming", summary="Streaming Image Generation")
-async def generate_image_streaming(request: ImageRequest = Body(...)):
-    prompt = request.prompt
-    num_inference_steps = request.num_inference_steps
-    guidance_scale = request.guidance_scale
-    format = request.format.lower()
-
-    def image_generator():
 
+async def load_pipeline():
+    """Loads the ONNX Stable Diffusion pipeline from Hugging Face Hub using snapshot_download."""
+    global pipeline
+    if pipeline is None:
         try:
-            start_time = time.time()
-            image = pipeline(
-                prompt,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale
-            ).images[0]
-            print(f"Image generation inference time: {time.time() - start_time:.3f}s")
-
-            img_byte_arr = io.BytesIO()
-            image_format = format.upper() if format in ["png", "jpeg"] else "PNG" # Default to PNG if format is invalid
-            image.save(img_byte_arr, format=image_format)
-            img_byte_arr = img_byte_arr.getvalue()
-            yield img_byte_arr
+            local_model_path = snapshot_download(
+                repo_id=repo_id,
+                local_dir=".", # Specify local_dir to ensure files are placed there
+                allow_patterns=["clip.opt/*.onnx"] # Specify necessary file patterns (adjust as needed)
+            )
+
+            pipeline = OnnxStableDiffusionPipeline.from_pretrained(
+                local_model_path,  # Use the local path from snapshot_download
+                provider="CPUExecutionProvider", # Or "CUDAExecutionProvider" if you have GPU
+            )
+            print(f"ONNX Stable Diffusion pipeline loaded successfully from {model_id} (ONNX revision) using snapshot_download from: {local_model_path}")
 
         except Exception as e:
-            print(f"Error processing image generation: {e}")
-            # yield error response? or just error out
-
-    media_type = f"image/{format}" if format in ["png", "jpeg"] else "image/png"
-    return StreamingResponse(
-        image_generator(),
-        media_type=media_type,
-        headers={"Cache-Control": "no-cache"},
-    )
-
-
-@app.post("/generate-image/full", summary="Full Image Generation")
-async def generate_image_full(request: ImageRequest = Body(...)):
-    prompt = request.prompt
-    num_inference_steps = request.num_inference_steps
-    guidance_scale = request.guidance_scale
-    format = request.format.lower()
-
-    start_time = time.time()
-    image = pipeline(
-        prompt,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale
-    ).images[0]
-    print(f"Full Image generation inference time: {time.time()-start_time:.3f}s")
+            print(f"Error loading ONNX pipeline using snapshot_download: {e}")
+            raise HTTPException(status_code=500, detail=f"Failed to load ONNX Stable Diffusion pipeline using snapshot_download: {e}")
+    return pipeline
 
+async def get_pipeline():
+    """Dependency to ensure pipeline is loaded before endpoint is called."""
+    return await load_pipeline()
 
-    img_byte_arr = io.BytesIO()
-    image_format = format.upper() if format in ["png", "jpeg"] else "PNG"
-    image.save(img_byte_arr, format=image_format)
-    img_byte_arr.seek(0)
+@app.on_event("startup")
+async def startup_event():
+    """Loads the pipeline on startup."""
+    await load_pipeline()
 
+@app.post("/generate-image/")
+async def generate_image(request: ImageRequest, pipeline_dep: OnnxStableDiffusionPipeline = Depends(get_pipeline)):
+    """
+    Generates an image based on the provided text prompt using the loaded ONNX Stable Diffusion pipeline.
+    """
+    try:
+        image = pipeline_dep(
+            request.prompt,
+            num_inference_steps=request.num_inference_steps,
+            guidance_scale=request.guidance_scale
+        ).images[0]
 
-    media_type = f"image/{format}" if format in ["png", "jpeg"] else "image/png"
-    return Response(content=img_byte_arr.read(), media_type=media_type)
-
-
-@app.get("/", response_class=HTMLResponse)
-def index():
-    return """
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>FastAPI Image Generation Demo</title>
-        <style>
-            body { font-family: Arial, sans-serif; }
-            .container { width: 80%; margin: auto; padding-top: 20px; }
-            h1 { text-align: center; }
-            .form-group { margin-bottom: 15px; }
-            label { display: block; margin-bottom: 5px; font-weight: bold; }
-            input[type="text"], input[type="number"], textarea, select { width: 100%; padding: 8px; box-sizing: border-box; margin-bottom: 10px; border: 1px solid #ccc; border-radius: 4px; }
-            textarea { height: 100px; }
-            button { padding: 10px 15px; border: none; color: white; background-color: #007bff; border-radius: 4px; cursor: pointer; }
-            button:hover { background-color: #0056b3; }
-            img { display: block; margin-top: 20px; max-width: 500px; } /* Adjust max-width as needed */
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <h1>FastAPI Image Generation Demo</h1>
-            <div class="form-group">
-                <label for="prompt">Text Prompt:</label>
-                <textarea id="prompt" rows="4" placeholder="Enter text prompt here"></textarea>
-            </div>
-            <div class="form-group">
-                <label for="num_inference_steps">Number of Inference Steps:</label>
-                <input type="number" id="num_inference_steps" value="50">
-            </div>
-            <div class="form-group">
-                <label for="guidance_scale">Guidance Scale:</label>
-                <input type="number" step="0.5" id="guidance_scale" value="7.5">
-            </div>
-            <div class="form-group">
-                <label for="format">Format:</label>
-                <select id="format">
-                    <option value="png" selected>PNG</option>
-                    <option value="jpeg">JPEG</option>
-                </select>
-            </div>
-            <div class="form-group">
-                <button onclick="generateStreamingImage()">Generate Streaming Image</button>
-                <button onclick="generateFullImage()">Generate Full Image</button>
-            </div>
-            <div id="image-container">
-                <img id="image" src="#" alt="Generated Image" style="display:none;">
-            </div>
-        </div>
-        <script>
-            function generateStreamingImage() {
-                const prompt = document.getElementById('prompt').value;
-                const num_inference_steps = document.getElementById('num_inference_steps').value;
-                const guidance_scale = document.getElementById('guidance_scale').value;
-                const format = document.getElementById('format').value;
-                const imageElement = document.getElementById('image');
-                const imageContainer = document.getElementById('image-container');
-
-                fetch('/generate-image/streaming', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json'
-                    },
-                    body: JSON.stringify({
-                        prompt: prompt,
-                        num_inference_steps: parseInt(num_inference_steps),
-                        guidance_scale: parseFloat(guidance_scale),
-                        format: format
-                    })
-                })
-                .then(response => response.blob())
-                .then(blob => {
-                    const imageUrl = URL.createObjectURL(blob);
-                    imageElement.src = imageUrl;
-                    imageElement.style.display = 'block'; // Show the image
-                    imageContainer.style.display = 'block'; // Show the container if hidden
-                });
-            }
-
-            function generateFullImage() {
-                const prompt = document.getElementById('prompt').value;
-                const num_inference_steps = document.getElementById('num_inference_steps').value;
-                const guidance_scale = document.getElementById('guidance_scale').value;
-                const format = document.getElementById('format').value;
-                const imageElement = document.getElementById('image');
-                const imageContainer = document.getElementById('image-container');
+        # Convert PIL Image to bytes for streaming response
+        img_byte_arr = BytesIO()
+        image.save(img_byte_arr, format="PNG")
+        img_byte_arr = img_byte_arr.getvalue()
 
+        return StreamingResponse(content=iter([img_byte_arr]), media_type="image/png")
 
-                fetch('/generate-image/full', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/json'
-                    },
-                    body: JSON.stringify({
-                        prompt: prompt,
-                        num_inference_steps: parseInt(num_inference_steps),
-                        guidance_scale: parseFloat(guidance_scale),
-                        format: format
-                    })
-                })
-                .then(response => response.blob())
-                .then(blob => {
-                    const imageUrl = URL.createObjectURL(blob);
-                    imageElement.src = imageUrl;
-                    imageElement.style.display = 'block'; // Show the image
-                    imageContainer.style.display = 'block'; // Show the container if hidden
-                });
-            }
-        </script>
-    </body>
-    </html>
-    """
+    except Exception as e:
+        print(f"Error during image generation: {e}")
+        raise HTTPException(status_code=500, detail=f"Image generation failed: {e}")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
+    uvicorn.run(app, host="0.0.0.0", port=8000, reload=True)