Update app.py

app.py

@@ -1,15 +1,211 @@
 
-# backend/app.py
+# # 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)
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
 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.responses import FileResponse
 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
@@ -19,31 +215,20 @@ 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_origins=[FRONTEND_URL] if FRONTEND_URL else ["*"],
+    allow_credentials=False,
     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 ------------------
@@ -72,6 +257,15 @@ 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:
@@ -87,41 +281,37 @@ def load_model(category: str, style: str):
     model = TransformerNet().to(device)
     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.")
+    return model
 
 # ------------------ 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()
-    ])
+def stylize_image(img: Image.Image, model):
     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)
+        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):
@@ -129,7 +319,7 @@ async def delete_file_after_delay(path: Path, delay: int = 180):
     try:
         if path.exists():
             path.unlink()
-            print(f"Deleted {path} after {delay}s")
+            print(f"Deleted {path}")
     except Exception as e:
         print(f"Error deleting file: {e}")
 
@@ -155,10 +345,15 @@ async def stylize(
 
     contents = await file.read()
     input_img = Image.open(io.BytesIO(contents)).convert("RGB")
-    output_tensor = stylize_image(input_img, model)
+
+    # 🔥 Run heavy task in background thread
+    output_tensor = await asyncio.to_thread(
+        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)
@@ -170,5 +365,5 @@ 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 or already deleted")
+        raise HTTPException(status_code=404, detail="File not found")
     return FileResponse(path, media_type="image/jpeg", filename=filename)