Update app.py

app.py

@@ -426,52 +426,44 @@ print("=" * 80)
 # ZeroGPU Tensor Preloading - Shared Components + Unique Transformers
 # =============================================================================
 
-print("Preloading models for ZeroGPU tensor packing...")
+print("Preloading all models for ZeroGPU tensor packing...")
 
-# ===== Load shared components from stage 1 =====
-print("  Loading shared components...")
-shared_ledger = pipeline.stage_1_model_ledger
-
-_video_encoder = shared_ledger.video_encoder()
-shared_ledger.video_encoder = lambda: _video_encoder
-
-_text_encoder = shared_ledger.text_encoder()
-shared_ledger.text_encoder = lambda: _text_encoder
-
-_embeddings_processor = shared_ledger.gemma_embeddings_processor()
-shared_ledger.gemma_embeddings_processor = lambda: _embeddings_processor
-
-_audio_decoder = shared_ledger.audio_decoder()
-shared_ledger.audio_decoder = lambda: _audio_decoder
-
-_vocoder = shared_ledger.vocoder()
-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
+# Get the actual loaded instances (these were already loaded lazily on first access)
+_transformer1 = pipeline.stage_1_model_ledger.transformer()
+_transformer2 = pipeline.stage_2_model_ledger.transformer()
+_video_encoder = pipeline.stage_1_model_ledger.video_encoder()
+_video_decoder = pipeline.stage_1_model_ledger.video_decoder()
+_audio_encoder = pipeline.stage_1_model_ledger.audio_encoder()
+_audio_decoder = pipeline.stage_1_model_ledger.audio_decoder()
+_vocoder = pipeline.stage_1_model_ledger.vocoder()
+_spatial_upsampler = pipeline.stage_1_model_ledger.spatial_upsampler()
+_text_encoder = pipeline.stage_1_model_ledger.text_encoder()
+_embeddings_processor = pipeline.stage_1_model_ledger.gemma_embeddings_processor()
+
+# Now replace ledger methods with lambdas returning cached instances
+# This prevents any accidental re-loading
+
+# Stage 1
+pipeline.stage_1_model_ledger.transformer = lambda: _transformer1
+pipeline.stage_1_model_ledger.video_encoder = lambda: _video_encoder
+pipeline.stage_1_model_ledger.video_decoder = lambda: _video_decoder
+pipeline.stage_1_model_ledger.audio_encoder = lambda: _audio_encoder
+pipeline.stage_1_model_ledger.audio_decoder = lambda: _audio_decoder
+pipeline.stage_1_model_ledger.vocoder = lambda: _vocoder
+pipeline.stage_1_model_ledger.spatial_upsampler = lambda: _spatial_upsampler
+pipeline.stage_1_model_ledger.text_encoder = lambda: _text_encoder
+pipeline.stage_1_model_ledger.gemma_embeddings_processor = lambda: _embeddings_processor
+
+# Stage 2 - point to same shared components
+pipeline.stage_2_model_ledger.transformer = lambda: _transformer2
 pipeline.stage_2_model_ledger.video_encoder = lambda: _video_encoder
-pipeline.stage_2_model_ledger.text_encoder = lambda: _text_encoder
-pipeline.stage_2_model_ledger.gemma_embeddings_processor = lambda: _embeddings_processor
+pipeline.stage_2_model_ledger.video_decoder = lambda: _video_decoder
+pipeline.stage_2_model_ledger.audio_encoder = lambda: _audio_encoder
 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")
-
-# ===== 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
-
-print("  Loading stage 2 transformer...")
-_transformer2 = pipeline.stage_2_model_ledger.transformer()
-pipeline.stage_2_model_ledger.transformer = lambda: _transformer2
+pipeline.stage_2_model_ledger.text_encoder = lambda: _text_encoder
+pipeline.stage_2_model_ledger.gemma_embeddings_processor = lambda: _embeddings_processor
 
 print("All models preloaded for ZeroGPU tensor packing!")
 print("=" * 80)
@@ -962,11 +954,12 @@ with gr.Blocks(title="LTX-2.3 Two-Stage HQ with LoRA Cache") as demo:
     generate_btn.click(
         fn=generate_video,
         inputs=[
-            prompt, negative_prompt, first_image, last_image, duration,
+            first_image, last_image, prompt, negative_prompt, duration, gpu_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, distilled_strength, pose_strength, general_strength, motion_strength,
+            gpu_duration,
+            distilled_strength, pose_strength, general_strength, motion_strength,
             dreamlay_strength, mself_strength, dramatic_strength, fluid_strength,
             liquid_strength, demopose_strength, voice_strength, realism_strength,
             transition_strength,