Update app.py

app.py

@@ -423,76 +423,62 @@ print("Pipeline initialized!")
 print("=" * 80)
 
 # =============================================================================
-# ZeroGPU Tensor Preloading - Adapted for Two Ledgers
+# ZeroGPU Tensor Preloading - Only Stage-Specific 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 all models for ZeroGPU tensor packing...")
-
-# ===== Stage 1 Model Ledger =====
-print("  Preloading Stage 1 models...")
-ledger1 = pipeline.stage_1_model_ledger
-
-_orig_transformer_factory1 = ledger1.transformer
-_orig_video_encoder_factory1 = ledger1.video_encoder
-_orig_video_decoder_factory1 = ledger1.video_decoder
-_orig_audio_decoder_factory1 = ledger1.audio_decoder
-_orig_vocoder_factory1 = ledger1.vocoder
-_orig_spatial_upsampler_factory1 = ledger1.spatial_upsampler
-_orig_text_encoder_factory1 = ledger1.text_encoder
-_orig_gemma_embeddings_factory1 = ledger1.gemma_embeddings_processor
-
-_transformer1 = _orig_transformer_factory1()
-_video_encoder1 = _orig_video_encoder_factory1()
-_video_decoder1 = _orig_video_decoder_factory1()
-_audio_decoder1 = _orig_audio_decoder_factory1()
-_vocoder1 = _orig_vocoder_factory1()
-_spatial_upsampler1 = _orig_spatial_upsampler_factory1()
-_text_encoder1 = _orig_text_encoder_factory1()
-_embeddings_processor1 = _orig_gemma_embeddings_factory1()
-
-ledger1.transformer = lambda: _transformer1
-ledger1.video_encoder = lambda: _video_encoder1
-ledger1.video_decoder = lambda: _video_decoder1
-ledger1.audio_decoder = lambda: _audio_decoder1
-ledger1.vocoder = lambda: _vocoder1
-ledger1.spatial_upsampler = lambda: _spatial_upsampler1
-ledger1.text_encoder = lambda: _text_encoder1
-ledger1.gemma_embeddings_processor = lambda: _embeddings_processor1
-
-print("    Stage 1 models preloaded")
-
-# ===== Stage 2 Model Ledger =====
-print("  Preloading Stage 2 models...")
-ledger2 = pipeline.stage_2_model_ledger
-
-_orig_transformer_factory2 = ledger2.transformer
-_orig_video_encoder_factory2 = ledger2.video_encoder
-_orig_video_decoder_factory2 = ledger2.video_decoder
-_orig_audio_decoder_factory2 = ledger2.audio_decoder
-_orig_vocoder_factory2 = ledger2.vocoder
-_orig_spatial_upsampler_factory2 = ledger2.spatial_upsampler
-_orig_text_encoder_factory2 = ledger2.text_encoder
-_orig_gemma_embeddings_factory2 = ledger2.gemma_embeddings_processor
-
-_transformer2 = _orig_transformer_factory2()
-_video_encoder2 = _orig_video_encoder_factory2()
-_video_decoder2 = _orig_video_decoder_factory2()
-_audio_decoder2 = _orig_audio_decoder_factory2()
-_vocoder2 = _orig_vocoder_factory2()
-_spatial_upsampler2 = _orig_spatial_upsampler_factory2()
-_text_encoder2 = _orig_text_encoder_factory2()
-_embeddings_processor2 = _orig_gemma_embeddings_factory2()
-
-ledger2.transformer = lambda: _transformer2
-ledger2.video_encoder = lambda: _video_encoder2
-ledger2.video_decoder = lambda: _video_decoder2
-ledger2.audio_decoder = lambda: _audio_decoder2
-ledger2.vocoder = lambda: _vocoder2
-ledger2.spatial_upsampler = lambda: _spatial_upsampler2
-ledger2.text_encoder = lambda: _text_encoder2
-ledger2.gemma_embeddings_processor = lambda: _embeddings_processor2
-
-print("    Stage 2 models preloaded")
+print("Preloading models for ZeroGPU tensor packing...")
+
+# ===== Preload SHARED components (load once, share) =====
+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
+
+_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
+
+# 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
+pipeline.stage_2_model_ledger.gemma_embeddings_processor = lambda: _embeddings_processor
+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
+
+print("    Shared components preloaded")
+
+# ===== Preload UNIQUE transformers per stage =====
+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
+
+# ===== 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)