| import torch |
| import traceback |
| import os |
| from PIL import Image |
| import io |
| import base64 |
| from fastapi import FastAPI, HTTPException, File, Form, UploadFile |
| from pydantic import BaseModel, Field |
| from diffusers import Flux2KleinPipeline |
| from fastapi.middleware.cors import CORSMiddleware |
| from typing import List |
|
|
| app = FastAPI() |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| print("Loading model...") |
| pipe = Flux2KleinPipeline.from_pretrained( |
| "black-forest-labs/FLUX.2-klein-base-4B", |
| torch_dtype=torch.bfloat16, |
| ) |
| pipe.to("cuda") |
| pipe.load_lora_weights("./sking_v73_flux_4b_000027000.safetensors") |
|
|
| print("Model loaded.") |
|
|
| class GenerateResponse(BaseModel): |
| images_base64: List[str] |
|
|
| @app.post("/api/generate", response_model=GenerateResponse) |
| async def generate(prompt: str=Form(...), guidance: float=Form(4.0), seed: int=Form(42), n_step: int = Form(100),file=File(...)): |
| images = [] |
| content = await file.read() |
| img = Image.open(io.BytesIO(content)).convert("RGBA") |
| images.append(img) |
| try: |
| pipeline_output = pipe( |
| image=images[0], |
| prompt="", |
| height=768, |
| width=768, |
| num_inference_steps=n_step, |
| guidance_scale=guidance, |
| num_images_per_prompt=1, |
| generator=torch.Generator("cuda").manual_seed(seed) |
| ) |
| |
| images = pipeline_output.images |
|
|
| b64_list = [] |
| for img in images: |
| buffered = io.BytesIO() |
| img.save(buffered, format="PNG") |
| img_str = base64.b64encode(buffered.getvalue()).decode() |
| b64_list.append(img_str) |
|
|
| return {"images_base64": b64_list} |
|
|
| except Exception as e: |
| traceback.print_exc() |
| print(f"Error: {e}") |
| raise HTTPException(status_code=500, detail=str(e)) |
|
|
| if __name__ == "__main__": |
| import uvicorn |
| uvicorn.run(app, host="0.0.0.0", port=10012) |
|
|
