Update app.py

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException, UploadFile, File
 from pydantic import BaseModel
 import requests
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
 import os
 import uuid
 from pathlib import Path
@@ -13,31 +14,40 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# Add CORS middleware to allow requests from Hugging Face Spaces frontend
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Adjust for production to specific origins
+    allow_origins=["*"],  # Adjust for production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# Mount static directory for serving files
+STATIC_DIR = Path("/home/user/app/static")
+STATIC_DIR.mkdir(parents=True, exist_ok=True)
+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
+
 # Define the request model
 class TryOnRequest(BaseModel):
     garmentDesc: str
     category: str
 
-# Base directory for file storage (use /tmp/gradio for Hugging Face Spaces)
-UPLOAD_DIR = Path("/tmp/gradio")
-UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+# Allowed file extensions
+ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png"}
 
 # Helper function to save file and generate public URL
 async def save_file_and_get_url(file: UploadFile) -> str:
     try:
+        # Validate file extension
+        file_extension = f".{file.filename.split('.')[-1].lower()}"
+        if file_extension not in ALLOWED_EXTENSIONS:
+            logger.error(f"Invalid file extension for {file.filename}: {file_extension}")
+            raise HTTPException(status_code=400, detail=f"File extension {file_extension} not allowed. Use JPG or PNG.")
+
         # Generate unique filename
-        file_extension = file.filename.split(".")[-1]
-        unique_filename = f"{uuid.uuid4()}.{file_extension}"
-        file_path = UPLOAD_DIR / unique_filename
+        unique_filename = f"{uuid.uuid4()}{file_extension}"
+        file_path = STATIC_DIR / unique_filename
         
         # Save file
         logger.info(f"Saving file to {file_path}")
@@ -51,8 +61,8 @@ async def save_file_and_get_url(file: UploadFile) -> str:
             raise HTTPException(status_code=500, detail="Failed to save file")
         
         # Generate public URL
-        # Use SPACE_ID environment variable or fallback to placeholder
-        public_url = f"https://tejani-tryapi.hf.space/file={str(file_path)}"
+        space_id = os.getenv("SPACE_ID", "tejani-tryapi")
+        public_url = f"https://{space_id}.hf.space/static/{unique_filename}"
         logger.info(f"Generated public URL: {public_url}")
         
         # Test URL accessibility
@@ -60,6 +70,8 @@ async def save_file_and_get_url(file: UploadFile) -> str:
             response = requests.head(public_url, timeout=5)
             if response.status_code != 200:
                 logger.warning(f"Public URL {public_url} returned status {response.status_code}")
+            else:
+                logger.info(f"Public URL {public_url} is accessible")
         except requests.exceptions.RequestException as e:
             logger.error(f"Failed to access public URL {public_url}: {str(e)}")
         
@@ -68,7 +80,7 @@ async def save_file_and_get_url(file: UploadFile) -> str:
         logger.error(f"Error in save_file_and_get_url: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error processing file: {str(e)}")
 
-# Updated endpoint to handle file uploads and proxy the request
+# Endpoint to handle file uploads and proxy the request
 @app.post("/try-on")
 async def try_on(
     human_img: UploadFile = File(...),
@@ -109,13 +121,13 @@ async def try_on(
             "garment_url": garment_url
         }
     except requests.exceptions.RequestException as e:
-        logger.error(f"Error forwarding request: {str(e)}")
+        logger.error(f"Error forwarding request: {unofficial: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error forwarding request: {str(e)}")
     except Exception as e:
         logger.error(f"Error in try_on endpoint: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
 
-# Health check endpoint for Hugging Face Spaces
+# Health check endpoint
 @app.get("/")
 async def root():
     return {"message": "FastAPI proxy for try-on API with file upload is running"}
@@ -124,8 +136,8 @@ async def root():
 @app.get("/list-files")
 async def list_files():
     try:
-        files = [str(f) for f in UPLOAD_DIR.glob("*") if f.is_file()]
-        logger.info(f"Files in {UPLOAD_DIR}: {files}")
+        files = [str(f) for f in STATIC_DIR.glob("*") if f.is_file()]
+        logger.info(f"Files in {STATIC_DIR}: {files}")
         return {"files": files}
     except Exception as e:
         logger.error(f"Error listing files: {str(e)}")