Update app.py

app.py

@@ -2,9 +2,11 @@ from fastapi import FastAPI, File, UploadFile, Request, HTTPException
 from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 import requests
-import time
 import asyncio
 from typing import Dict
+import os
+import json
+import re
 
 app = FastAPI()
 
@@ -385,7 +387,7 @@ HTML_CONTENT = """
             <div class="loading-spinner" id="loadingSpinner"></div>
         </form>
         <div class="result-container" id="resultContainer"></div>
-                <div class="file-types">
+        <div class="file-types">
             Allowed file types: .zip, .mp4, .txt, .mp3, all image types, .pdf
         </div>
     </div>
@@ -409,7 +411,7 @@ HTML_CONTENT = """
         const progressContainer = document.getElementById('progressContainer');
         const loadingSpinner = document.getElementById('loadingSpinner');
         const resultContainer = document.getElementById('resultContainer');
-                const dropZone = document.getElementById('dropZone');
+        const dropZone = document.getElementById('dropZone');
         const modal = document.getElementById('embedModal');
         const span = document.getElementsByClassName("close")[0];
         const embedLinkInput = document.getElementById('embedLink');
@@ -417,14 +419,14 @@ HTML_CONTENT = """
 
         fileInput.addEventListener('change', handleFileSelect);
 
-        uploadForm.addEventListener('submit', async (e) => {
+        uploadForm.addEventListener('submit', (e) => {
             e.preventDefault();
             if (fileInput.files.length > 0) {
-                await uploadFile(fileInput.files[0]);
+                uploadFile(fileInput.files[0]);
             }
         });
 
-        dropZone.addEventListener('dragover', (e) => {
+                dropZone.addEventListener('dragover', (e) => {
             e.preventDefault();
             dropZone.classList.add('drag-over');
         });
@@ -486,43 +488,47 @@ HTML_CONTENT = """
             const formData = new FormData();
             formData.append('file', file);
 
-            let uploadedBytes = 0;
-            const chunkSize = 1024 * 1024; // 1MB chunks
-            const totalSize = file.size;
-
-            while (uploadedBytes < totalSize) {
-                const chunk = file.slice(uploadedBytes, uploadedBytes + chunkSize);
-                formData.set('file', chunk, file.name);
-
+            while (true) {
                 try {
-                    const response = await fetch('/upload', {
-                        method: 'POST',
-                        body: formData,
-                        headers: {
-                            'Content-Range': `bytes ${uploadedBytes}-${uploadedBytes + chunk.size - 1}/${totalSize}`
+                    const xhr = new XMLHttpRequest();
+                    xhr.open('POST', '/upload', true);
+                    xhr.upload.onprogress = (event) => updateProgress(event, progressBar.querySelector('.progress'));
+                    
+                    xhr.onload = function() {
+                        if (xhr.status === 200) {
+                            const response = JSON.parse(xhr.responseText);
+                            if (response.url) {
+                                addResultLink(response.url, file.name);
+                                resetUploadState();
+                                return;
+                            } else {
+                                throw new Error('Upload failed: ' + response.error);
+                            }
+                        } else {
+                            throw new Error(`HTTP error! status: ${xhr.status}`);
                         }
-                    });
+                    };
 
-                    if (!response.ok) {
-                        throw new Error(`HTTP error! status: ${response.status}`);
-                    }
+                    xhr.onerror = function() {
+                        throw new Error('Network error occurred');
+                    };
 
-                    const result = await response.json();
-                    if (result.url) {
-                        addResultLink(result.url, file.name);
-                        break;
-                    }
+                    xhr.send(formData);
 
-                    uploadedBytes += chunk.size;
-                    updateProgress(uploadedBytes, totalSize, progressBar.querySelector('.progress'));
+                    // Wait for the request to complete
+                    await new Promise((resolve, reject) => {
+                        xhr.onloadend = resolve;
+                        xhr.onerror = reject;
+                    });
+
+                    break; // Success, exit the loop
                 } catch (error) {
                     console.error('Upload error:', error);
-                    await new Promise(resolve => setTimeout(resolve, 1000)); // Wait 1 second before retrying
-                    // The loop will continue from where it left off
+                    // Wait for a short time before retrying
+                    await new Promise(resolve => setTimeout(resolve, 1000));
+                    // The loop will continue, retrying the upload
                 }
             }
-
-            resetUploadState();
         }
 
         function createProgressBar(fileName) {
@@ -541,9 +547,11 @@ HTML_CONTENT = """
             return container;
         }
 
-        function updateProgress(uploadedBytes, totalBytes, progressBar) {
-            const percentComplete = (uploadedBytes / totalBytes) * 100;
-            progressBar.style.width = percentComplete + '%';
+        function updateProgress(event, progressBar) {
+            if (event.lengthComputable) {
+                const percentComplete = (event.loaded / event.total) * 100;
+                progressBar.style.width = percentComplete + '%';
+            }
         }
 
         function resetUploadState() {
@@ -621,37 +629,85 @@ async def handle_upload(request: Request, file: UploadFile = File(...)):
     if not content_range:
         raise HTTPException(status_code=400, detail="Content-Range header is missing")
 
-    _, range_str = content_range.split(' ')
-    start_byte, end_byte, total_size = map(int, range_str.replace('bytes', '').split('-')[0].split('/'))
+    # Parse Content-Range header
+    content_range_match = re.match(r'bytes (\d+)-(\d+)/(\d+)', content_range)
+    if not content_range_match:
+        raise HTTPException(status_code=400, detail="Invalid Content-Range header format")
+
+    start_byte = int(content_range_match.group(1))
+    end_byte = int(content_range_match.group(2))
+    total_size = int(content_range_match.group(3))
+
+    if start_byte > end_byte or end_byte >= total_size:
+        raise HTTPException(status_code=400, detail="Invalid Content-Range header values")
 
     if not file:
         raise HTTPException(status_code=400, detail="No file part")
     if file.filename == '':
         raise HTTPException(status_code=400, detail="No selected file")
 
-    cookies = await get_cookies()
-    if 'csrftoken' not in cookies or 'sessionid' not in cookies:
-        raise HTTPException(status_code=500, detail="Failed to obtain necessary cookies")
+    # Create temporary directory to store the chunks if it doesn't exist
+    upload_dir = os.path.join('/tmp', 'uploads')
+    os.makedirs(upload_dir, exist_ok=True)
+
+    # Temporary file path
+    temp_file_path = os.path.join(upload_dir, file.filename)
+
+    # Path to store upload metadata
+    meta_file_path = temp_file_path + '.json'
 
     if start_byte == 0:
+        # Start of a new upload
+        cookies = await get_cookies()
+        if 'csrftoken' not in cookies or 'sessionid' not in cookies:
+            raise HTTPException(status_code=500, detail="Failed to obtain necessary cookies")
+
+        # Initiate the upload
         upload_result = await initiate_upload(cookies, file.filename, file.content_type)
         if not upload_result or 'upload_url' not in upload_result:
             raise HTTPException(status_code=500, detail="Failed to initiate upload")
+        
+        # Save the upload_url and serving_url to a metadata file
+        with open(meta_file_path, 'w') as meta_file:
+            json.dump(upload_result, meta_file)
     else:
-        upload_result = {'upload_url': file.filename}  # Use filename as a key to retrieve the upload URL
+        # For subsequent chunks, read the metadata file to get upload_url and serving_url
+        if not os.path.exists(meta_file_path):
+            raise HTTPException(status_code=400, detail="Upload metadata not found")
+        with open(meta_file_path, 'r') as meta_file:
+            upload_result = json.load(meta_file)
 
+    # Read the chunk content
     file_content = await file.read()
-    upload_success = await upload_chunk(upload_result['upload_url'], file_content, start_byte, end_byte, total_size)
-    if not upload_success:
-        raise HTTPException(status_code=500, detail="Chunk upload failed")
 
-    if end_byte + 1 >= total_size:
+    # Write the chunk to the temporary file at the correct offset
+    with open(temp_file_path, 'ab') as f:
+        f.seek(start_byte)
+        f.write(file_content)
+
+    # Check if the upload is complete
+    file_size = os.path.getsize(temp_file_path)
+    if file_size == total_size:
+        # Read the entire file content
+        with open(temp_file_path, 'rb') as f:
+            full_file_content = f.read()
+
+        # Upload the file to Replicate
+        upload_success = await retry_upload(upload_result['upload_url'], full_file_content, file.content_type)
+        if not upload_success:
+            raise HTTPException(status_code=500, detail="Failed to upload file to Replicate")
+
         original_url = upload_result.get('serving_url')
         if original_url:
             mirrored_url = f"/rbxg/{original_url.split('/pbxt/')[1]}"
+            # Clean up the temporary files
+            os.remove(temp_file_path)
+            os.remove(meta_file_path)
             return JSONResponse(content={"url": mirrored_url})
-    
-    return JSONResponse(content={"status": "chunk uploaded"})
+        else:
+            raise HTTPException(status_code=500, detail="serving_url not found")
+    else:
+        return JSONResponse(content={"status": "chunk uploaded"})
 
 @app.get("/rbxg/{path:path}")
 async def handle_video_stream(path: str, request: Request):
@@ -718,7 +774,7 @@ async def embed_video(url: str, thumbnail: str):
 async def get_cookies() -> Dict[str, str]:
     try:
         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'
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'
         })
         return dict(response.cookies)
     except Exception as e:
@@ -728,9 +784,9 @@ async def get_cookies() -> Dict[str, str]:
 async def initiate_upload(cookies: Dict[str, str], filename: str, content_type: str) -> Dict:
     url = f'https://replicate.com/api/upload/{filename}?content_type={content_type}'
     try:
-        response = requests.post(url, cookies=cookies, headers={
+        headers = {
             'X-CSRFToken': cookies.get('csrftoken'),
-            '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',
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)',
             'Referer': 'https://replicate.com/levelsio/neon-tokyo',
             'Origin': 'https://replicate.com',
             'Accept': '*/*',
@@ -738,48 +794,34 @@ async def initiate_upload(cookies: Dict[str, str], filename: str, content_type:
             'Accept-Encoding': 'identity',
             'Sec-Fetch-Dest': 'empty',
             'Sec-Fetch-Mode': 'cors',
-            'Sec-Fetch-Site': 'same-origin',
-            'Sec-GPC': '1',
-            'Priority': 'u=1, i'
-        })
+            'Sec-Fetch-Site': 'same-origin'
+        }
+        response = requests.post(url, cookies=cookies, headers=headers)
         return response.json()
     except Exception as e:
         print(f'Error initiating upload: {e}')
         raise
 
-async def upload_chunk(upload_url: str, chunk_content: bytes, start_byte: int, end_byte: int, total_size: int) -> bool:
+async def upload_file(upload_url: str, file_content: bytes, content_type: str) -> bool:
     try:
         headers = {
-            'Content-Type': 'application/octet-stream',
-            'Content-Range': f'bytes {start_byte}-{end_byte}/{total_size}'
+            'Content-Type': content_type
         }
-        response = requests.put(upload_url, data=chunk_content, headers=headers)
+        response = requests.put(upload_url, data=file_content, headers=headers)
         return response.status_code in [200, 201, 204]
     except Exception as e:
-        print(f'Error uploading chunk: {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:
     retries = 0
     while retries < max_retries:
-        try:
-            success = await upload_file(upload_url, file_content, content_type)
-            if success:
-                return True
+        success = await upload_file(upload_url, file_content, content_type)
+        if success:
+            return True
+        else:
             print(f"Upload attempt {retries + 1} failed. Retrying...")
-        except Exception as e:
-            print(f"Error during upload attempt {retries + 1}: {e}")
-        
-        retries += 1
-        await asyncio.sleep(delay)
-        delay = min(delay * 2, 60)  # Exponential backoff, capped at 60 seconds
-    
-    return False
-
-async def upload_file(upload_url: str, file_content: bytes, content_type: str) -> bool:
-    try:
-        response = requests.put(upload_url, data=file_content, headers={'Content-Type': content_type})
-        return response.status_code == 200
-    except Exception as e:
-        print(f'Error uploading file: {e}')
-        return False
+            retries += 1
+            await asyncio.sleep(delay)
+            delay = min(delay * 2, 60)  # Exponential backoff
+    return False