File size: 4,093 Bytes
cfeb09d 2002864 cfeb09d 2002864 ac80bb1 cfeb09d 2002864 cfeb09d 2002864 8d5cd96 2002864 cfeb09d 2002864 cfeb09d 2002864 cfeb09d 2002864 cfeb09d 2002864 cfeb09d 2002864 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import torch
import logging
import io
from pathlib import Path
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from diffusers import StableDiffusionPipeline
from peft import PeftModel
from fastapi.middleware.cors import CORSMiddleware
# ==================================
# Setup: App, Logging, Constants
# ==================================
app = FastAPI(
title="Stable Diffusion LoRA Inference API",
description="The Flake API to generate images using a Stable Diffusion model fine-tuned with LoRA.",
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Or specify your frontend's URL
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
)
# Constants for model paths
BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5"
LORA_UNET_PATH = "./lora_unet"
LORA_TEXT_ENCODER_PATH = "./lora_text"
# Global variable to hold the pipeline
pipe = None
# ==================================
# Pydantic Models for API
# ==================================
class ImageRequest(BaseModel):
"""Request model for image generation."""
prompt: str
num_steps: int = 30
guidance_scale: float = 7.5
# ==================================
# Model Loading & Application Events
# ==================================
@app.on_event("startup")
def load_model():
"""
Load the model and adapters once when the application starts.
This is efficient as the model stays in memory.
"""
global pipe
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
logging.info(f"Using device: {device}")
# Load the base pipeline
logging.info(f"Loading base model: {BASE_MODEL_ID}")
pipe = StableDiffusionPipeline.from_pretrained(
BASE_MODEL_ID,
torch_dtype=torch_dtype,
safety_checker=None,
)
# Load and fuse the LoRA adapters
try:
logging.info(f"Loading and attaching UNet LoRA from: {LORA_UNET_PATH}")
pipe.unet = PeftModel.from_pretrained(pipe.unet, LORA_UNET_PATH)
logging.info(f"Loading and attaching Text Encoder LoRA from: {LORA_TEXT_ENCODER_PATH}")
pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, LORA_TEXT_ENCODER_PATH)
except Exception as e:
logging.error(f"Fatal: Could not load LoRA adapters. Server will not work. Error: {e}")
# In a real scenario, you might want the app to fail startup if models don't load.
raise RuntimeError(f"Failed to load LoRA adapters: {e}") from e
pipe.to(device)
logging.info("Model and LoRA adapters loaded successfully.")
# ==================================
# API Endpoints
# ==================================
@app.get("/")
def read_root():
"""Root endpoint to check if the API is running."""
return {"status": "API is running", "docs_url": "/docs"}
@app.post("/generate-image")
async def generate_image_endpoint(request: ImageRequest):
"""
Endpoint to generate an image.
Receives a POST request with a JSON body containing the prompt.
Returns the generated image as a PNG file.
"""
if pipe is None:
raise HTTPException(status_code=503, detail="Model is not loaded or failed to load.")
try:
logging.info(f"Generating image for prompt: '{request.prompt}'")
image = pipe(
prompt=request.prompt,
num_inference_steps=request.num_steps,
guidance_scale=request.guidance_scale,
).images[0]
# Save image to an in-memory buffer
buffer = io.BytesIO()
image.save(buffer, format="PNG")
buffer.seek(0)
# Return the image as a streaming response
return StreamingResponse(buffer, media_type="image/png")
except Exception as e:
logging.error(f"Error during image generation: {e}")
raise HTTPException(status_code=500, detail=str(e)) |