Update app.py

app.py

@@ -148,12 +148,12 @@ def update_image_input(mode):
     else:
         return gr.update(visible=False)
 
-def wait_for_model_with_retry(replicate_client, model_name, max_retries=5, initial_wait=10):
+def wait_for_model_with_retry(model_name, max_retries=5, initial_wait=10):
     """Wait for model to be ready with retry logic"""
     for attempt in range(max_retries):
         try:
             # Try to get model info
-            model = replicate_client.models.get(model_name)
+            model = replicate.models.get(model_name)
             return True
         except Exception as e:
             if attempt < max_retries - 1:
@@ -212,14 +212,11 @@ def generate_video(mode, prompt, image, aspect_ratio, seed, api_key_input, progr
         
         progress(0.2, desc="Checking model availability...")
         
-        # Create Replicate client with extended timeout
-        client = replicate.Client(
-            api_token=token,
-            timeout=300  # 5 minutes timeout
-        )
+        # Set up Replicate with the API token
+        replicate.api_token = token
         
         # Wait for model to be ready
-        model_ready = wait_for_model_with_retry(client, "bytedance/seedance-1-lite")
+        model_ready = wait_for_model_with_retry("bytedance/seedance-1-lite")
         if not model_ready:
             return None, "⏳ Model is still booting up. Please try again in a few minutes."
         
@@ -227,54 +224,47 @@ def generate_video(mode, prompt, image, aspect_ratio, seed, api_key_input, progr
         
         # Run Replicate with retry logic
         max_attempts = 3
+        output = None
+        
         for attempt in range(max_attempts):
             try:
-                # Use prediction API for better control
-                prediction = client.predictions.create(
-                    version="bytedance/seedance-1-lite:latest",
-                    input=input_params,
-                    webhook_completed=None
-                )
-                
-                # Poll for completion with extended timeout
+                # Run Replicate - use the model directly without version specifier
                 start_time = time.time()
-                timeout_seconds = 300  # 5 minutes
-                
-                while prediction.status not in ["succeeded", "failed", "canceled"]:
-                    if time.time() - start_time > timeout_seconds:
-                        return None, "⏱️ Generation timed out. The server might be under heavy load. Please try again."
-                    
-                    time.sleep(2)  # Poll every 2 seconds
-                    prediction.reload()
-                    
-                    # Update progress
-                    elapsed = time.time() - start_time
-                    progress_val = min(0.3 + (elapsed / timeout_seconds) * 0.4, 0.7)
-                    progress(progress_val, desc=f"Generating video... ({int(elapsed)}s)")
                 
-                if prediction.status == "failed":
-                    error_msg = getattr(prediction, 'error', 'Unknown error')
-                    if "cold boot" in str(error_msg).lower() or "starting" in str(error_msg).lower():
-                        if attempt < max_attempts - 1:
-                            progress(0.3, desc=f"Model is starting up, retrying... (Attempt {attempt + 2}/{max_attempts})")
-                            time.sleep(30)  # Wait 30 seconds before retry
-                            continue
-                    return None, f"❌ Generation failed: {error_msg}"
+                # Create a generator for progress tracking
+                output_generator = replicate.run(
+                    "bytedance/seedance-1-lite",
+                    input=input_params
+                )
                 
-                elif prediction.status == "canceled":
-                    return None, "❌ Generation was canceled."
+                # If it's a generator, iterate through it
+                if hasattr(output_generator, '__iter__') and not isinstance(output_generator, (str, bytes)):
+                    for event in output_generator:
+                        elapsed = time.time() - start_time
+                        progress_val = min(0.3 + (elapsed / 300) * 0.4, 0.7)
+                        progress(progress_val, desc=f"Generating video... ({int(elapsed)}s)")
+                        output = event  # Keep the last event as output
+                else:
+                    # If it's not a generator, use it directly
+                    output = output_generator
                 
-                # Success - get output
-                output = prediction.output
-                break
+                # If we got output, break the retry loop
+                if output:
+                    break
                 
             except replicate.exceptions.ReplicateError as e:
-                if "timeout" in str(e).lower() and attempt < max_attempts - 1:
+                error_str = str(e)
+                if "cold boot" in error_str.lower() or "starting" in error_str.lower():
+                    if attempt < max_attempts - 1:
+                        progress(0.3, desc=f"Model is starting up, retrying... (Attempt {attempt + 2}/{max_attempts})")
+                        time.sleep(30)
+                        continue
+                elif "timeout" in error_str.lower() and attempt < max_attempts - 1:
                     progress(0.3, desc=f"Timeout occurred, retrying... (Attempt {attempt + 2}/{max_attempts})")
                     time.sleep(10)
                     continue
                 else:
-                    return None, f"❌ Replicate API error: {str(e)}"
+                    return None, f"❌ Replicate API error: {error_str}"
             except Exception as e:
                 if attempt < max_attempts - 1:
                     progress(0.3, desc=f"Error occurred, retrying... (Attempt {attempt + 2}/{max_attempts})")
@@ -283,6 +273,10 @@ def generate_video(mode, prompt, image, aspect_ratio, seed, api_key_input, progr
                 else:
                     return None, f"❌ Unexpected error: {str(e)}"
         
+        # Check if we got output
+        if not output:
+            return None, "❌ Failed to generate video after multiple attempts."
+        
         progress(0.7, desc="Downloading video...")
         
         # Get video data