Upload 4 files

app.py

@@ -0,0 +1,329 @@
+
+# # backend/main.py
+# import os, io, uuid, sys, json, asyncio
+# from pathlib import Path
+# from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, BackgroundTasks
+# from fastapi.responses import FileResponse, JSONResponse
+# from fastapi.middleware.cors import CORSMiddleware
+# from fastapi.staticfiles import StaticFiles
+# from PIL import Image
+# import torch
+# from torchvision import transforms
+
+
+
+
+# BASE_DIR = Path(__file__).resolve().parent
+# sys.path.append(str(BASE_DIR / "helpers"))
+# from helpers.transform_net import TransformerNet
+
+# app = FastAPI()
+
+# # -------- CORS: add your frontend origin (dev: http://localhost:5173) ----------
+# FRONTEND_URL = os.environ.get("FRONTEND_URL", "http://localhost:5173")
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=[FRONTEND_URL],
+#     allow_credentials=True,
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
+
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# # ---------- outputs dir ------------
+# OUTPUT_DIR = BASE_DIR / "outputs"
+# OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+# # mount static so /download/<file> serves images directly
+# app.mount("/download", StaticFiles(directory=str(OUTPUT_DIR)), name="download")
+
+# # ---------- load models.json + your model cache (keep your existing logic) ----------
+# models_json_path = BASE_DIR / "models.json"
+# if not models_json_path.exists():
+#     raise RuntimeError(f"models.json not found at {models_json_path}")
+
+# with open(models_json_path, "r") as f:
+#     MODEL_PATHS = json.load(f)
+
+# # convert to absolute paths (your existing code)
+# for cat, styles in MODEL_PATHS.items():
+#     for style_name, rel_path in styles.items():
+#         p = Path(rel_path)
+#         if not p.is_absolute():
+#             MODEL_PATHS[cat][style_name] = str((BASE_DIR / rel_path).resolve())
+
+# models = {}
+# def load_model(category: str, style: str):
+#     key = (category, style)
+#     if key in models:
+#         return models[key]
+#     if category not in MODEL_PATHS or style not in MODEL_PATHS[category]:
+#         raise HTTPException(status_code=400, detail="Invalid category/style")
+#     path = MODEL_PATHS[category][style]
+#     if not os.path.exists(path):
+#         raise HTTPException(status_code=404, detail=f"Model file not found: {path}")
+#     model = TransformerNet().to(device)
+#     model.load_state_dict(torch.load(path, map_location=device))
+#     model.eval()
+#     models[key] = model
+#     return model
+
+# def save_image_tensor(tensor, path: Path):
+#     img = tensor.detach().float().cpu()[0].clamp(0,1).permute(1,2,0).numpy() * 255
+#     Image.fromarray(img.astype("uint8")).save(path)
+
+# def stylize_image(img: Image.Image, model, img_size: int = 512):
+#     transform = transforms.Compose([transforms.Resize(img_size), transforms.ToTensor()])
+#     x = transform(img).unsqueeze(0).to(device)
+#     with torch.no_grad(), torch.cuda.amp.autocast(enabled=True):
+#         y = model(x)
+#     return y
+
+# # cleanup helper
+# async def delete_file_after_delay(path: Path, delay: int = 180):
+#     await asyncio.sleep(delay)
+#     try:
+#         if path.exists():
+#             path.unlink()
+#             print(f"🧹 Deleted {path} after {delay} sec")
+#     except Exception as e:
+#         print("⚠️ error deleting file:", e)
+
+# # --------- routes ----------
+
+
+# @app.get("/")
+# async def root():
+#     return {"message": "Backend is running!"}
+
+
+# # # -------------------- API Routes --------------------
+# @app.get("/api/styles")
+# async def get_styles():
+#     return MODEL_PATHS
+
+
+
+# @app.post("/api/stylize")
+# async def stylize(
+#     request: Request,
+#     background_tasks: BackgroundTasks,
+#     file: UploadFile = File(...),
+#     category: str = Form(...),
+#     style: str = Form(...),
+# ):
+#     model = load_model(category, style)
+#     contents = await file.read()
+#     input_img = Image.open(io.BytesIO(contents)).convert("RGB")
+#     output_tensor = stylize_image(input_img, model)
+#     filename = f"{uuid.uuid4().hex}.jpg"
+#     out_path = OUTPUT_DIR / filename
+#     save_image_tensor(output_tensor, out_path)
+
+#     # schedule deletion
+#     background_tasks.add_task(delete_file_after_delay, out_path, 180)
+
+#     # build absolute URL using the request base URL (works in dev and production)
+#     base = str(request.base_url).rstrip('/')
+#     image_url = f"{base}/download/{filename}"
+#     return {"image_url": image_url}
+
+# # (OPTIONAL) keep a file endpoint if you want; not necessary because StaticFiles serves it:
+# @app.get("/api/download/{filename}")
+# async def download(filename: str):
+#     path = OUTPUT_DIR / filename
+#     return FileResponse(path, media_type="image/jpeg", filename=filename)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+# backend/main.py
+import os, io, uuid, sys, json, asyncio
+from pathlib import Path
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request, BackgroundTasks
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from PIL import Image
+import torch
+from torchvision import transforms
+
+# ------------------ BASE SETUP ------------------
+
+BASE_DIR = Path(__file__).resolve().parent
+sys.path.append(str(BASE_DIR / "helpers"))
+from helpers.transform_net import TransformerNet
+
+app = FastAPI()
+
+# ------------------ CORS ------------------
+# In HF Spaces dashboard, set environment variable:
+#   FRONTEND_URL = https://your-app.vercel.app
+# For local dev it defaults to localhost:5173
+
+FRONTEND_URL = os.environ.get("FRONTEND_URL", "http://localhost:5173")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[FRONTEND_URL],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ------------------ DEVICE ------------------
+# HF Spaces free tier = CPU only
+# cuda.amp.autocast is disabled on CPU to avoid warnings
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+use_amp = device.type == "cuda"
+print(f"Running on: {device}")
+
+# ------------------ OUTPUTS ------------------
+
+OUTPUT_DIR = BASE_DIR / "outputs"
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+app.mount("/download", StaticFiles(directory=str(OUTPUT_DIR)), name="download")
+
+# ------------------ MODELS ------------------
+
+models_json_path = BASE_DIR / "models.json"
+if not models_json_path.exists():
+    raise RuntimeError(f"models.json not found at {models_json_path}")
+
+with open(models_json_path, "r") as f:
+    MODEL_PATHS = json.load(f)
+
+# Convert relative paths to absolute
+for cat, styles in MODEL_PATHS.items():
+    for style_name, rel_path in styles.items():
+        p = Path(rel_path)
+        if not p.is_absolute():
+            MODEL_PATHS[cat][style_name] = str((BASE_DIR / rel_path).resolve())
+
+# In-memory model cache
+models = {}
+
+def load_model(category: str, style: str):
+    key = (category, style)
+    if key in models:
+        return models[key]
+
+    if category not in MODEL_PATHS or style not in MODEL_PATHS[category]:
+        raise HTTPException(status_code=400, detail="Invalid category/style")
+
+    path = MODEL_PATHS[category][style]
+    if not os.path.exists(path):
+        raise HTTPException(status_code=404, detail=f"Model file not found: {path}")
+
+    model = TransformerNet().to(device)
+    model.load_state_dict(torch.load(path, map_location=device))
+    model.eval()
+    models[key] = model
+    print(f"Loaded model: {category}/{style}")
+    return model
+
+# Preload all models at startup
+# Since each model is only 10-11 MB, all fit easily in 16 GB free RAM
+@app.on_event("startup")
+async def preload_all_models():
+    print("Preloading all models...")
+    for cat, styles in MODEL_PATHS.items():
+        for style in styles:
+            try:
+                load_model(cat, style)
+            except Exception as e:
+                print(f"Warning: Could not load {cat}/{style} — {e}")
+    print(f"Done. {len(models)} model(s) loaded.")
+
+# ------------------ IMAGE UTILS ------------------
+
+def save_image_tensor(tensor, path: Path):
+    img = tensor.detach().float().cpu()[0].clamp(0, 1).permute(1, 2, 0).numpy() * 255
+    Image.fromarray(img.astype("uint8")).save(path)
+
+def stylize_image(img: Image.Image, model, img_size: int = 512):
+    transform = transforms.Compose([
+        transforms.Resize(img_size),
+        transforms.ToTensor()
+    ])
+    x = transform(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        # autocast only when GPU is available, safe no-op on CPU
+        with torch.cuda.amp.autocast(enabled=use_amp):
+            y = model(x)
+    return y
+
+# ------------------ CLEANUP ------------------
+
+async def delete_file_after_delay(path: Path, delay: int = 180):
+    await asyncio.sleep(delay)
+    try:
+        if path.exists():
+            path.unlink()
+            print(f"Deleted {path} after {delay}s")
+    except Exception as e:
+        print(f"Error deleting file: {e}")
+
+# ------------------ ROUTES ------------------
+
+@app.get("/")
+async def root():
+    return {"message": "Backend is running!", "device": str(device)}
+
+@app.get("/api/styles")
+async def get_styles():
+    return MODEL_PATHS
+
+@app.post("/api/stylize")
+async def stylize(
+    request: Request,
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(...),
+    category: str = Form(...),
+    style: str = Form(...),
+):
+    model = load_model(category, style)
+
+    contents = await file.read()
+    input_img = Image.open(io.BytesIO(contents)).convert("RGB")
+    output_tensor = stylize_image(input_img, model)
+
+    filename = f"{uuid.uuid4().hex}.jpg"
+    out_path = OUTPUT_DIR / filename
+    save_image_tensor(output_tensor, out_path)
+
+    background_tasks.add_task(delete_file_after_delay, out_path, 180)
+
+    base_url = str(request.base_url).rstrip("/")
+    return {"image_url": f"{base_url}/download/{filename}"}
+
+@app.get("/api/download/{filename}")
+async def download(filename: str):
+    path = OUTPUT_DIR / filename
+    if not path.exists():
+        raise HTTPException(status_code=404, detail="File not found or already deleted")
+    return FileResponse(path, media_type="image/jpeg", filename=filename)

models.json

@@ -0,0 +1,26 @@
+{
+  "human-face": {
+    "van-gogh": "./models/van-gogh-humans-epoch15.pth",
+    "claude-monet": "./models/claude-monet-humans-epoch10.pth"
+  },
+
+  "cars": {
+    "van-gogh": "./models/car-van-gogh-gpt_edgeAMP_epoch25.pth",
+    "claude-monet": "./models/car-claude-monet-gpt_edgeAMP_epoch12.pth"
+  },
+
+  "cats": {
+    "van-gogh": "./models/cats-van-gogh-gpt_edgeAMP_epoch12.pth",
+    "claude-monet": "./models/cat-claude-monet-gpt_edgeAMP_epoch12.pth"    
+  },
+
+  "dogs": {
+    "van-gogh": "./models/dog-van-gogh-gpt_edgeAMP_epoch12.pth",
+    "claude-monet": "./models/dog-claude-monet-gpt_edgeAMP_epoch12.pth"    
+  },
+
+  "landscape": {
+    "van-gogh": "./models/landscape-van-gogh-gpt_edgeAMP_epoch30.pth",
+    "claude-monet": "./models/landscape-claude-monet-gpt_edgeAMP_epoch30.pth"
+  }
+}

package-lock.json

@@ -0,0 +1,6 @@
+{
+  "name": "backend",
+  "lockfileVersion": 3,
+  "requires": true,
+  "packages": {}
+}

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+pillow
+torch
+torchvision
+python-multipart
+aiofiles