Update app.py

app.py

@@ -34,7 +34,7 @@ except subprocess.CalledProcessError as e:
     print(f"Error installing package: {e.stderr}")
     sys.exit(1)
 
-# --- 3. Refactor the demo script using a robust line-by-line patch ---
+# --- 3. Refactor the demo script using a robust state-machine patcher ---
 demo_script_path = Path("demo/gradio_demo.py")
 print(f"Refactoring {demo_script_path} for ZeroGPU lazy loading...")
 
@@ -63,13 +63,16 @@ try:
             print("Model and processor loaded successfully on GPU worker.")
 """
     
-    # --- Perform the line-by-line modifications ---
+    # --- Perform the line-by-line modifications using a state machine ---
     new_lines = []
     # Add 'import spaces' at the top if it doesn't exist
     if not any("import spaces" in line for line in lines):
         new_lines.append("import spaces\n")
 
-    patched = False
+    # State machine variables
+    in_generate_method = False
+    patched_generate_method = False
+    
     for line in lines:
         # Defer the initial model loading to prevent PicklingError
         if "self.load_model()" in line and "def __init__" in "".join(lines[lines.index(line)-2:lines.index(line)]):
@@ -77,21 +80,30 @@ try:
             new_lines.append("        self.model = None\n")
             new_lines.append("        self.processor = None\n")
             print("Successfully patched __init__ to prevent startup model load.")
-        # Find the generation method to add the decorator and lazy-loading logic
-        elif "def generate_podcast_streaming(self," in line:
+        
+        # Start of the target method
+        elif "def generate_podcast_streaming(self," in line and not patched_generate_method:
             new_lines.append("    @spaces.GPU(duration=120)\n")
             new_lines.append(line)
-        elif "-> Iterator[tuple]:" in line and "generate_podcast_streaming" in new_lines[-1]:
+            in_generate_method = True
+        
+        # End of the target method signature
+        elif "-> Iterator[tuple]:" in line and in_generate_method:
             new_lines.append(line)
-            # Indent the lazy load code correctly
+            # Indent and insert the lazy load code
             for code_line in lazy_load_code.strip().split('\n'):
                 new_lines.append(' ' * 8 + code_line + '\n')
-            patched = True
+            
+            # Reset state
+            in_generate_method = False
+            patched_generate_method = True
             print("Successfully patched generation method for lazy loading.")
+        
+        # All other lines
         else:
             new_lines.append(line)
             
-    if not patched:
+    if not patched_generate_method:
         print("\033[91mError: Failed to apply the lazy-loading patch. The target method signature may have changed.\033[0m")
         sys.exit(1)