Update app.py

app.py

@@ -423,20 +423,15 @@ print("Pipeline initialized!")
 print("=" * 80)
 
 # =============================================================================
-# ZeroGPU Tensor Preloading - Only Stage-Specific Transformers
+# ZeroGPU Tensor Preloading - Shared Components + Unique Transformers
 # =============================================================================
-# Many components are SHARED between stage 1 and stage 2:
-# - video_encoder, text_encoder (Gemma), audio_encoder, vocoder, embeddings_processor
-# These only need to be loaded ONCE.
-# Only the TRANSFORMER is different per stage (different LoRA weights).
 
 print("Preloading models for ZeroGPU tensor packing...")
 
-# ===== Preload SHARED components (load once, share) =====
+# ===== Load shared components from stage 1 =====
 print("  Loading shared components...")
 shared_ledger = pipeline.stage_1_model_ledger
 
-# These are shared - load once from stage 1
 _video_encoder = shared_ledger.video_encoder()
 shared_ledger.video_encoder = lambda: _video_encoder
 
@@ -455,6 +450,9 @@ shared_ledger.vocoder = lambda: _vocoder
 _spatial_upsampler = shared_ledger.spatial_upsampler()
 shared_ledger.spatial_upsampler = lambda: _spatial_upsampler
 
+_video_decoder = shared_ledger.video_decoder()
+shared_ledger.video_decoder = lambda: _video_decoder
+
 # Point stage 2 to same shared components
 pipeline.stage_2_model_ledger.video_encoder = lambda: _video_encoder
 pipeline.stage_2_model_ledger.text_encoder = lambda: _text_encoder
@@ -462,10 +460,11 @@ pipeline.stage_2_model_ledger.gemma_embeddings_processor = lambda: _embeddings_p
 pipeline.stage_2_model_ledger.audio_decoder = lambda: _audio_decoder
 pipeline.stage_2_model_ledger.vocoder = lambda: _vocoder
 pipeline.stage_2_model_ledger.spatial_upsampler = lambda: _spatial_upsampler
+pipeline.stage_2_model_ledger.video_decoder = lambda: _video_decoder
 
-print("    Shared components preloaded")
+print("Shared components preloaded")
 
-# ===== Preload UNIQUE transformers per stage =====
+# ===== Load ONLY stage-specific transformers =====
 print("  Loading stage 1 transformer...")
 _transformer1 = pipeline.stage_1_model_ledger.transformer()
 pipeline.stage_1_model_ledger.transformer = lambda: _transformer1
@@ -474,12 +473,6 @@ print("  Loading stage 2 transformer...")
 _transformer2 = pipeline.stage_2_model_ledger.transformer()
 pipeline.stage_2_model_ledger.transformer = lambda: _transformer2
 
-# ===== Preload video_decoder (may be shared, but preload for safety) =====
-print("  Loading video decoder...")
-_video_decoder1 = shared_ledger.video_decoder()
-shared_ledger.video_decoder = lambda: _video_decoder1
-pipeline.stage_2_model_ledger.video_decoder = lambda: _video_decoder1
-
 print("All models preloaded for ZeroGPU tensor packing!")
 print("=" * 80)
 print("Pipeline ready!")
@@ -643,11 +636,25 @@ def on_highres_toggle(first_image, last_image, high_res):
 
 
 def get_duration(
-    prompt, negative_prompt, first_image, last_image,
-    duration, seed, randomize_seed, height, width, enhance_prompt,
-    video_cfg_scale, video_stg_scale, video_rescale_scale, video_a2v_scale,
-    audio_cfg_scale, audio_stg_scale, audio_rescale_scale, audio_v2a_scale,
-    gpu_duration,
+    first_image,
+    last_image,
+    prompt: str,
+    negative_prompt: str,
+    duration: float,
+    seed: int,
+    randomize_seed: bool,
+    height: int,
+    width: int,
+    enhance_prompt: bool,
+    video_cfg_scale: float,
+    video_stg_scale: float,
+    video_rescale_scale: float,
+    video_a2v_scale: float,
+    audio_cfg_scale: float,
+    audio_stg_scale: float,
+    audio_rescale_scale: float,
+    audio_v2a_scale: float,
+    gpu_duration: float,
     progress=None,
 ) -> int:
     return int(gpu_duration)
@@ -656,10 +663,10 @@ def get_duration(
 @spaces.GPU(duration=get_duration)
 @torch.inference_mode()
 def generate_video(
-    prompt: str,
-    negative_prompt: str,
     first_image,
     last_image,
+    prompt: str,
+    negative_prompt: str,
     duration: float,
     seed: int,
     randomize_seed: bool,
@@ -754,21 +761,6 @@ def generate_video(
 
         output_path = tempfile.mktemp(suffix=".mp4")
         encode_video(
-            video=video,
-            fps=DEFAULT_FRAME_RATE,
-            audio=audio,
-            output_path=output_path,
-            video_chunks_number=video_chunks_number,
-        )
-
-        log_memory("after encode_video")
-        return str(output_path), current_seed
-
-    except Exception as e:
-        import traceback
-        log_memory("on error")
-        print(f"Error: {str(e)}\n{traceback.format_exc()}")
-        return None, current_seed
 
 
 # =============================================================================