| import os, io, base64, threading, traceback |
| import torch |
| import numpy as np |
| from PIL import Image, ImageDraw |
| from fastapi import FastAPI, UploadFile, File |
| from fastapi.responses import HTMLResponse, JSONResponse |
| import uvicorn |
|
|
| app = FastAPI() |
| MODEL_LOADED = False |
| LOAD_ERROR = "" |
| pipe = None |
|
|
| def load_model(): |
| global pipe, MODEL_LOADED, LOAD_ERROR |
| try: |
| print("📥 Loading model on CPU...") |
| from diffusers import StableDiffusionInpaintPipeline |
| pipe = StableDiffusionInpaintPipeline.from_pretrained( |
| "runwayml/stable-diffusion-inpainting", |
| torch_dtype=torch.float32, |
| safety_checker=None, |
| requires_safety_checker=False, |
| ) |
| |
| pipe.enable_attention_slicing() |
| MODEL_LOADED = True |
| print("✅ Model ready on CPU!") |
| except Exception as e: |
| LOAD_ERROR = str(e) |
| print(f"❌ {e}") |
|
|
| threading.Thread(target=load_model, daemon=True).start() |
|
|
| def pil_to_b64(img): |
| buf = io.BytesIO() |
| img.save(buf, format="PNG") |
| return base64.b64encode(buf.getvalue()).decode() |
|
|
| def make_mask(size): |
| w, h = size |
| mask = Image.new("L", size, 0) |
| draw = ImageDraw.Draw(mask) |
| draw.rectangle([w*0.05, h*0.18, w*0.95, h*0.68], fill=255) |
| draw.rectangle([w*0.0, h*0.18, w*0.15, h*0.58], fill=255) |
| draw.rectangle([w*0.85, h*0.18, w*1.0, h*0.58], fill=255) |
| return mask.convert("RGB") |
|
|
| @app.get("/", response_class=HTMLResponse) |
| async def index(): |
| return HTMLResponse(open("/app/index.html").read()) |
|
|
| @app.get("/status") |
| async def status(): |
| return {"loaded": MODEL_LOADED, "error": LOAD_ERROR} |
|
|
| @app.post("/tryon") |
| async def tryon(person: UploadFile = File(...), garment: UploadFile = File(...)): |
| if not MODEL_LOADED: |
| return JSONResponse({"status":"loading","message":"Model still loading, please wait and retry."}, status_code=503) |
| try: |
| SIZE = (512, 768) |
| person_img = Image.open(io.BytesIO(await person.read())).convert("RGB").resize(SIZE) |
| mask_img = make_mask(SIZE) |
|
|
| prompt = "Person wearing a clean stylish garment, photorealistic, high quality fashion photo, same pose, same background" |
| negative = "nude, deformed, blurry, bad anatomy, extra limbs, watermark, logo, text, disfigured" |
|
|
| result = pipe( |
| prompt=prompt, |
| negative_prompt=negative, |
| image=person_img, |
| mask_image=mask_img, |
| height=SIZE[1], |
| width=SIZE[0], |
| num_inference_steps=25, |
| guidance_scale=7.5, |
| strength=0.95, |
| ).images[0] |
|
|
| return JSONResponse({"status":"ok","image": pil_to_b64(result)}) |
| except Exception as e: |
| traceback.print_exc() |
| return JSONResponse({"status":"error","message":str(e)}, status_code=500) |
|
|
| if __name__ == "__main__": |
| uvicorn.run(app, host="0.0.0.0", port=7860) |
|
|
