Update app.py

app.py

@@ -80,6 +80,13 @@ class DownloadResponse(BaseModel):
     filename: str = None
     download_url: str = None
 
+
+
+def verify_file_size(file_path: str, expected_size: int):
+    actual_size = os.path.getsize(file_path)
+    if actual_size < expected_size:
+        raise Exception(f"Downloaded file size mismatch: expected {expected_size} bytes, got {actual_size} bytes")
+
 @app.post("/download", response_model=DownloadResponse)
 async def download_file(request: DownloadRequest):
     try:
@@ -128,6 +135,9 @@ async def download_file(request: DownloadRequest):
             
         downloaded_file = files[0]
         logger.info(f"Downloaded file: {downloaded_file}")
+
+        # Verify the size of the downloaded file (expected size in bytes, e.g., 7000000 for 7MB)
+        verify_file_size(downloaded_file, expected_size=7000000)
         
         # Generate the download URL
         space_url = os.getenv("SPACE_URL", "http://localhost:7860")
@@ -161,6 +171,7 @@ async def download_file(request: DownloadRequest):
         if 'original_dir' in locals():
             os.chdir(original_dir)
 
+
 @app.get("/test")
 async def test():
     """Test endpoint to verify setup"""