| from typing import Literal, Optional, TYPE_CHECKING |
|
|
| import numpy as np |
| from fastapi import FastAPI |
| from fastapi.responses import StreamingResponse, JSONResponse |
| from pydantic import BaseModel, Field |
| from platform import system |
|
|
| if TYPE_CHECKING: |
| from flux_pipeline import FluxPipeline |
|
|
| if system() == "Windows": |
| MAX_RAND = 2**16 - 1 |
| else: |
| MAX_RAND = 2**32 - 1 |
|
|
|
|
| class AppState: |
| model: "FluxPipeline" |
|
|
|
|
| class FastAPIApp(FastAPI): |
| state: AppState |
|
|
|
|
| class LoraArgs(BaseModel): |
| scale: Optional[float] = 1.0 |
| path: Optional[str] = None |
| name: Optional[str] = None |
| action: Optional[Literal["load", "unload"]] = "load" |
|
|
|
|
| class LoraLoadResponse(BaseModel): |
| status: Literal["success", "error"] |
| message: Optional[str] = None |
|
|
|
|
| class GenerateArgs(BaseModel): |
| prompt: str |
| width: Optional[int] = Field(default=720) |
| height: Optional[int] = Field(default=1024) |
| num_steps: Optional[int] = Field(default=24) |
| guidance: Optional[float] = Field(default=3.5) |
| seed: Optional[int] = Field( |
| default_factory=lambda: np.random.randint(0, MAX_RAND), gt=0, lt=MAX_RAND |
| ) |
| strength: Optional[float] = 1.0 |
| init_image: Optional[str] = None |
|
|
|
|
| app = FastAPIApp() |
|
|
|
|
| @app.post("/generate") |
| def generate(args: GenerateArgs): |
| """ |
| Generates an image from the Flux flow transformer. |
| |
| Args: |
| args (GenerateArgs): Arguments for image generation: |
| |
| - `prompt`: The prompt used for image generation. |
| |
| - `width`: The width of the image. |
| |
| - `height`: The height of the image. |
| |
| - `num_steps`: The number of steps for the image generation. |
| |
| - `guidance`: The guidance for image generation, represents the |
| influence of the prompt on the image generation. |
| |
| - `seed`: The seed for the image generation. |
| |
| - `strength`: strength for image generation, 0.0 - 1.0. |
| Represents the percent of diffusion steps to run, |
| setting the init_image as the noised latent at the |
| given number of steps. |
| |
| - `init_image`: Base64 encoded image or path to image to use as the init image. |
| |
| Returns: |
| StreamingResponse: The generated image as streaming jpeg bytes. |
| """ |
| result = app.state.model.generate(**args.model_dump()) |
| return StreamingResponse(result, media_type="image/jpeg") |
|
|
|
|
| @app.post("/lora", response_model=LoraLoadResponse) |
| def lora_action(args: LoraArgs): |
| """ |
| Loads or unloads a LoRA checkpoint into / from the Flux flow transformer. |
| |
| Args: |
| args (LoraArgs): Arguments for the LoRA action: |
| |
| - `scale`: The scaling factor for the LoRA weights. |
| - `path`: The path to the LoRA checkpoint. |
| - `name`: The name of the LoRA checkpoint. |
| - `action`: The action to perform, either "load" or "unload". |
| |
| Returns: |
| LoraLoadResponse: The status of the LoRA action. |
| """ |
| try: |
| if args.action == "load": |
| app.state.model.load_lora(args.path, args.scale, args.name) |
| elif args.action == "unload": |
| app.state.model.unload_lora(args.name if args.name else args.path) |
| else: |
| return JSONResponse( |
| content={ |
| "status": "error", |
| "message": f"Invalid action, expected 'load' or 'unload', got {args.action}", |
| }, |
| status_code=400, |
| ) |
| except Exception as e: |
| return JSONResponse( |
| status_code=500, content={"status": "error", "message": str(e)} |
| ) |
| return JSONResponse(status_code=200, content={"status": "success"}) |
|
|
