Upload 3 files

Dockerfile

@@ -0,0 +1,12 @@
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y ffmpeg git && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,111 @@
+from fastapi import FastAPI, UploadFile, File, Form
+import uuid, os
+from PIL import Image
+
+from rembg import remove
+from realesrgan import RealESRGAN
+from gfpgan import GFPGANer
+from diffusers import StableDiffusionImg2ImgPipeline
+import torch
+
+app = FastAPI()
+
+UPLOAD_DIR = "uploads"
+OUTPUT_DIR = "outputs"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ===== Load models ONCE =====
+
+# Real-ESRGAN
+upscaler = RealESRGAN(torch.device(DEVICE), scale=4)
+upscaler.load_weights("RealESRGAN_x4.pth")
+
+# GFPGAN
+gfpgan = GFPGANer(
+    model_path=None,
+    upscale=1,
+    arch="clean",
+    channel_multiplier=2,
+    device=DEVICE
+)
+
+# Stable Diffusion Img2Img
+sd_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16
+).to(DEVICE)
+
+# ===== Endpoints =====
+
+@app.post("/remove-bg")
+async def remove_bg(file: UploadFile = File(...)):
+    img = Image.open(file.file).convert("RGBA")
+    out = remove(img)
+
+    fname = f"{uuid.uuid4().hex}.png"
+    path = os.path.join(OUTPUT_DIR, fname)
+    out.save(path)
+
+    return {"file": fname}
+
+
+@app.post("/enhance")
+async def enhance(file: UploadFile = File(...)):
+    img = Image.open(file.file).convert("RGB")
+    out = upscaler.predict(img)
+
+    fname = f"{uuid.uuid4().hex}.png"
+    out.save(os.path.join(OUTPUT_DIR, fname))
+    return {"file": fname}
+
+
+@app.post("/upscale")
+async def upscale(file: UploadFile = File(...), scale: int = Form(2)):
+    img = Image.open(file.file).convert("RGB")
+    upscaler.scale = scale
+    out = upscaler.predict(img)
+
+    fname = f"{uuid.uuid4().hex}.png"
+    out.save(os.path.join(OUTPUT_DIR, fname))
+    return {"file": fname}
+
+
+@app.post("/retouch")
+async def retouch(file: UploadFile = File(...)):
+    img = Image.open(file.file).convert("RGB")
+    _, _, out = gfpgan.enhance(
+        img,
+        has_aligned=False,
+        only_center_face=False,
+        paste_back=True
+    )
+
+    fname = f"{uuid.uuid4().hex}.png"
+    Image.fromarray(out).save(os.path.join(OUTPUT_DIR, fname))
+    return {"file": fname}
+
+
+@app.post("/edit")
+async def edit_image(
+    file: UploadFile = File(...),
+    prompt: str = Form(...),
+    strength: float = Form(0.6)
+):
+    img = Image.open(file.file).convert("RGB").resize((512, 512))
+
+    result = sd_pipe(
+        prompt=prompt,
+        image=img,
+        strength=strength,
+        guidance_scale=7.5
+    ).images[0]
+
+    fname = f"{uuid.uuid4().hex}.png"
+    result.save(os.path.join(OUTPUT_DIR, fname))
+    return {
+        "prompt": prompt,
+        "file": fname
+    }

requirements.txt

@@ -0,0 +1,14 @@
+fastapi
+uvicorn
+torch
+torchvision
+diffusers
+transformers
+accelerate
+opencv-python
+pillow
+numpy
+scipy
+rembg
+realesrgan
+gfpgan