infer.py

infer.py

@@ -0,0 +1,71 @@
+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)