| 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 |
|
|
|
|
|
|
|
|
|
|
| |
| |
| |
| 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=["*"], |
| allow_credentials=True, |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| logging.basicConfig( |
| level=logging.INFO, |
| format="%(asctime)s - %(levelname)s - %(message)s", |
| ) |
|
|
| |
| BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5" |
| LORA_UNET_PATH = "./lora_unet" |
| LORA_TEXT_ENCODER_PATH = "./lora_text" |
|
|
| |
| pipe = None |
|
|
| |
| |
| |
| class ImageRequest(BaseModel): |
| """Request model for image generation.""" |
| prompt: str |
| num_steps: int = 30 |
| guidance_scale: float = 7.5 |
|
|
| |
| |
| |
| @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}") |
|
|
| |
| logging.info(f"Loading base model: {BASE_MODEL_ID}") |
| pipe = StableDiffusionPipeline.from_pretrained( |
| BASE_MODEL_ID, |
| torch_dtype=torch_dtype, |
| safety_checker=None, |
| ) |
| |
| |
| 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}") |
| |
| raise RuntimeError(f"Failed to load LoRA adapters: {e}") from e |
|
|
| pipe.to(device) |
| logging.info("Model and LoRA adapters loaded successfully.") |
|
|
|
|
| |
| |
| |
| @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] |
| |
| |
| buffer = io.BytesIO() |
| image.save(buffer, format="PNG") |
| buffer.seek(0) |
|
|
| |
| 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)) |
