Update app.py

app.py

@@ -184,26 +184,36 @@ print("=" * 80)
 print("Preloading all models for ZeroGPU tensor packing...")
 print("This may take a few minutes...")
 
-# Access all pipeline components to force tensor loading
-# These are accessed but stored as lambda closures so ZeroGPU packs them
-_transformer = pipeline.stage_1.transformer
+# Access pipeline components to force tensor loading
+# TI2VidTwoStagesHQPipeline exposes these as class attributes (not through a ledger)
 _video_encoder = pipeline.prompt_encoder.video_encoder
 _video_decoder = pipeline.video_decoder
 _audio_decoder = pipeline.audio_decoder
 _spatial_upsampler = pipeline.upsampler
+_prompt_encoder = pipeline.prompt_encoder
 
-# Trigger actual loading by calling them once
-# This forces tensors into memory so ZeroGPU can pack them
-print("  Loading transformer...")
-_transformer()
+# Trigger actual loading by accessing the underlying models
+# These are lazily loaded on first access
 print("  Loading video encoder...")
-_video_encoder()
+_video_encoder.model
 print("  Loading video decoder...")
-_video_decoder()
+_video_decoder.model
 print("  Loading audio decoder...")
-_audio_decoder()
+_audio_decoder.model
 print("  Loading spatial upsampler...")
-_spatial_upsampler()
+_spatial_upsampler.model
+print("  Loading prompt encoder (Gemma)...")
+_prompt_encoder.gemma
+
+# Trigger diffusion stages to load their transformer weights
+# We do this by creating a small dummy inference to force loading
+print("  Loading stage 1 transformer...")
+with torch.no_grad():
+    dummy_latent = torch.randn(1, 16, 9, 64, 64, device='cpu', dtype=torch.bfloat16)
+    dummy_sigmas = torch.tensor([1.0, 0.5, 0.0], device='cpu', dtype=torch.float32)
+    # Access the internal transformer through stage_1._transformer_ctx
+    _ = pipeline.stage_1._transformer_ctx
+    _ = pipeline.stage_2._transformer_ctx
 
 print("All models preloaded for ZeroGPU tensor packing!")
 print("=" * 80)