Update app.py

app.py

@@ -1,7 +1,6 @@
 from fastapi import FastAPI, File, UploadFile, Request
-from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse, FileResponse
+from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 import requests
-import time
 import asyncio
 from typing import Dict
 import os
@@ -445,7 +444,7 @@ HTML_CONTENT = """
             text-align: center;
         }
 
-                .quick-open-content img,
+        .quick-open-content img,
         .quick-open-content video,
         .quick-open-content audio {
             max-width: 100%;
@@ -889,7 +888,7 @@ HTML_CONTENT = """
             historyModal.style.display = "block";
         }
 
-                function quickOpen(url, fileName, originalExtension) {
+        function quickOpen(url, fileName, originalExtension) {
             quickOpenContent.innerHTML = '';
             const fullUrl = window.location.origin + url;
             
@@ -947,7 +946,7 @@ async def handle_upload(file: UploadFile = File(...)):
 
     cookies = await get_cookies()
     if 'csrftoken' not in cookies or 'sessionid' not in cookies:
-        return JSONResponse(content={"error": "Failed"}, status_code=500)
+        return JSONResponse(content={"error": "Failed to get cookies from Replicate."}, status_code=500)
 
     original_extension = os.path.splitext(file.filename)[1][1:]
     supported_types = ['mp4', 'png', 'jpg', 'jpeg', 'gif', 'mp3', 'pdf', 'txt']
@@ -961,12 +960,15 @@ async def handle_upload(file: UploadFile = File(...)):
 
     upload_result = await initiate_upload(cookies, temp_filename, content_type)
     if not upload_result or 'upload_url' not in upload_result:
-        return JSONResponse(content={"error": "Failed to upload"}, status_code=500)
+        return JSONResponse(content={"error": "Failed to initiate upload with Replicate."}, status_code=500)
+
+    # Log the real Replicate URL
+    print(f"Real Replicate URL: {upload_result['serving_url']}")
 
     file_content = await file.read()
     upload_success = await retry_upload(upload_result['upload_url'], file_content, content_type)
     if not upload_success:
-        return JSONResponse(content={"error": "FAILED GOD MAN AFTER alot of attempts"}, status_code=500)
+        return JSONResponse(content={"error": "File upload to Replicate failed after multiple attempts."}, status_code=500)
 
     original_url = upload_result['serving_url']
     mirrored_url = f"/upload/{original_url.split('/pbxt/')[1]}"
@@ -985,25 +987,31 @@ async def handle_file_stream(path: str, request: Request):
     
     range_header = request.headers.get('Range')
     headers = {'Range': range_header} if range_header else {}
-    response = requests.get(original_url, headers=headers, stream=True)
+    
+    try:
+        response = requests.get(original_url, headers=headers, stream=True)
+        response.raise_for_status()
 
-    def generate():
-        for chunk in response.iter_content(chunk_size=8192):
-            yield chunk
+        def generate():
+            for chunk in response.iter_content(chunk_size=8192):
+                yield chunk
 
-    headers = dict(response.headers)
-    headers['Access-Control-Allow-Origin'] = '*'
-    headers['Content-Disposition'] = 'inline'
+        response_headers = dict(response.headers)
+        response_headers['Access-Control-Allow-Origin'] = '*'
+        response_headers['Content-Disposition'] = 'inline'
 
-    if response.status_code == 206:
-        headers['Content-Range'] = response.headers.get('Content-Range')
+        if response.status_code == 206:
+            response_headers['Content-Range'] = response.headers.get('Content-Range')
 
-    original_extension = os.path.splitext(path)[1][1:]
-    content_type, _ = mimetypes.guess_type(f"file.{original_extension}")
-    if content_type:
-        headers['Content-Type'] = content_type
+        original_extension = os.path.splitext(path)[1][1:]
+        content_type, _ = mimetypes.guess_type(f"file.{original_extension}")
+        if content_type:
+            response_headers['Content-Type'] = content_type
+
+        return StreamingResponse(generate(), status_code=response.status_code, headers=response_headers)
+    except requests.exceptions.RequestException as e:
+        return JSONResponse(content={"error": f"Failed to fetch file from Replicate: {e}"}, status_code=500)
 
-    return StreamingResponse(generate(), status_code=response.status_code, headers=headers)
 
 @app.get("/embed")
 async def embed_video(url: str, thumbnail: str):
@@ -1051,8 +1059,9 @@ async def get_cookies() -> Dict[str, str]:
         response = requests.get('https://replicate.com/levelsio/neon-tokyo', headers={
             'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36'
         })
+        response.raise_for_status()
         return dict(response.cookies)
-    except Exception as e:
+    except requests.exceptions.RequestException as e:
         print(f'Error fetching the page: {e}')
         return {}
 
@@ -1066,37 +1075,37 @@ async def initiate_upload(cookies: Dict[str, str], filename: str, content_type:
             'Origin': 'https://replicate.com',
             'Accept': '*/*',
             'Accept-Language': 'en-US,en;q=0.5',
-            'Accept-Encoding': 'identity',
+            'Accept-Encoding': 'gzip, deflate, br',
             'Sec-Fetch-Dest': 'empty',
             'Sec-Fetch-Mode': 'cors',
             'Sec-Fetch-Site': 'same-origin',
-            'Sec-GPC': '1',
-            'Priority': 'u=1, i'
         })
+        response.raise_for_status()
         return response.json()
-    except Exception as e:
+    except requests.exceptions.RequestException as e:
         print(f'Error initiating upload: {e}')
-        raise
+        return {}
 
-async def upload_file(upload_url: str, file_content: bytes, content_type: str) -> bool:
+async def upload_file_to_gcs(upload_url: str, file_content: bytes, content_type: str) -> bool:
     try:
         response = requests.put(upload_url, data=file_content, headers={'Content-Type': content_type})
+        response.raise_for_status()
         return response.status_code == 200
-    except Exception as e:
+    except requests.exceptions.RequestException as e:
         print(f'Error uploading file: {e}')
         return False
 
 async def retry_upload(upload_url: str, file_content: bytes, content_type: str, max_retries: int = 5, delay: int = 1) -> bool:
-    while True:
+    for attempt in range(max_retries):
         try:
-            success = await upload_file(upload_url, file_content, content_type)
+            success = await upload_file_to_gcs(upload_url, file_content, content_type)
             if success:
                 return True
-            print("Upload failed. Retrying...")
+            print(f"Upload attempt {attempt + 1} failed. Retrying in {delay} seconds...")
         except Exception as e:
-            print(f"Error during upload: {e}")
+            print(f"Error during upload attempt {attempt + 1}: {e}")
         
         await asyncio.sleep(delay)
         delay = min(delay * 2, 60)
     
-    return False
+    return False