Update app.py

app.py

@@ -1,211 +1,14 @@
-
-# # backend/app.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")
-# # FRONTEND_URL = "https://image-stylizer-deploy.vercel.app"
-
-# app.add_middleware(
-#     CORSMiddleware,
-#     allow_origins=[
-#         FRONTEND_URL
-#         # "https://image-stylizer-deploy.vercel.app",
-#         # "http://localhost:5173",  # for local testing
-#     ],
-#     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()
-
-#     model = torch.jit.script(model)
-    
-#     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 = 256):
-#     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
-#         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)
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+# backend/app.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
+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
 
-# ------------------ PERFORMANCE SETTINGS ------------------
-
-torch.set_num_threads(1)  # 🔥 critical for HF CPU
-
 # ------------------ BASE SETUP ------------------
 
 BASE_DIR = Path(__file__).resolve().parent
@@ -215,20 +18,31 @@ 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")
+# FRONTEND_URL = "https://image-stylizer-deploy.vercel.app"
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[FRONTEND_URL] if FRONTEND_URL else ["*"],
-    allow_credentials=False,
+    allow_origins=[
+        FRONTEND_URL
+        # "https://image-stylizer-deploy.vercel.app",
+        # "http://localhost:5173",  # for local testing
+    ],
+    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 ------------------
@@ -257,15 +71,6 @@ for cat, styles in MODEL_PATHS.items():
 # In-memory model cache
 models = {}
 
-# ------------------ GLOBAL TRANSFORM ------------------
-
-transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.ToTensor()
-])
-
-# ------------------ MODEL LOADER ------------------
-
 def load_model(category: str, style: str):
     key = (category, style)
     if key in models:
@@ -282,36 +87,42 @@ def load_model(category: str, style: str):
     model.load_state_dict(torch.load(path, map_location=device))
     model.eval()
 
-    # 🔥 TorchScript optimization
     model = torch.jit.script(model)
-
-    # 🔥 Warmup (removes first-request delay)
-    dummy = torch.randn(1, 3, 256, 256).to(device)
-    with torch.no_grad():
-        model(dummy)
-
+    
     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 stylize_image(img: Image.Image, 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 = 256):
+    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
         y = model(x)
     return y
 
-def save_image_tensor(tensor, path: Path):
-    img = tensor.detach().cpu()[0].clamp(0, 1).permute(1, 2, 0).numpy() * 255
-    Image.fromarray(img.astype("uint8")).save(
-        path,
-        format="JPEG",
-        quality=85,
-        optimize=True
-    )
-
 # ------------------ CLEANUP ------------------
 
 async def delete_file_after_delay(path: Path, delay: int = 180):
@@ -319,7 +130,7 @@ async def delete_file_after_delay(path: Path, delay: int = 180):
     try:
         if path.exists():
             path.unlink()
-            print(f"Deleted {path}")
+            print(f"Deleted {path} after {delay}s")
     except Exception as e:
         print(f"Error deleting file: {e}")
 
@@ -345,15 +156,10 @@ async def stylize(
 
     contents = await file.read()
     input_img = Image.open(io.BytesIO(contents)).convert("RGB")
-
-    # 🔥 Run heavy task in background thread
-    output_tensor = await asyncio.to_thread(
-        stylize_image, input_img, model
-    )
+    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)
@@ -365,5 +171,6 @@ async def stylize(
 async def download(filename: str):
     path = OUTPUT_DIR / filename
     if not path.exists():
-        raise HTTPException(status_code=404, detail="File not found")
-    return FileResponse(path, media_type="image/jpeg", filename=filename)
+        raise HTTPException(status_code=404, detail="File not found or already deleted")
+    return FileResponse(path, media_type="image/jpeg", filename=filename)
+