Update app.py

app.py

@@ -181,65 +181,55 @@ print("=" * 80)
 print("Preloading all models for ZeroGPU tensor packing...")
 print("This may take a few minutes...")
 
-# The TI2VidTwoStagesHQPipeline uses context managers for lazy loading.
-# We need to enter the contexts, capture the loaded models, AND preserve them
-# by replacing the pipeline's internal references with lambdas that hold them.
-# This prevents garbage collection and allows ZeroGPU to pack them.
+# TI2VidTwoStagesHQPipeline uses:
+# - Builder methods that return models directly when called
+# - Context methods that return context managers when called
+# We need to call these methods, capture the results, and preserve them
 
-# 1. Load transformer via _transformer_ctx (enter context to load, store result)
+# 1. Load transformer via _transformer_ctx() (call first, then enter)
 print("  Loading stage 1 transformer...")
-pipeline.stage_1._transformer_ctx.__enter__()
-# Capture the actual model from the context
-_stage_1_transformer = pipeline.stage_1._transformer_ctx.__dict__.get('transformer') or \
+_ctx = pipeline.stage_1._transformer_ctx()  # Get context manager
+_ctx.__enter__()  # Enter context
+_stage_1_transformer = _ctx.__dict__.get('transformer') or \
                        getattr(pipeline.stage_1, '_transformer', None)
-# Replace _transformer_ctx with lambda that returns the captured model
-pipeline.stage_1._transformer_ctx = type('ctx', (), {
-    '__enter__': lambda s: _stage_1_transformer,
-    '__exit__': lambda s, *a: None,
-    '__call__': lambda s, *a, **kw: _stage_1_transformer(*a, **kw)
-})()
+# Replace _transformer_ctx with a lambda that returns cached model
+pipeline.stage_1._transformer_ctx = lambda: _ctx
 print(f"    Captured stage 1 transformer: {type(_stage_1_transformer)}")
 
 print("  Loading stage 2 transformer...")
-pipeline.stage_2._transformer_ctx.__enter__()
-_stage_2_transformer = pipeline.stage_2._transformer_ctx.__dict__.get('transformer') or \
+_ctx = pipeline.stage_2._transformer_ctx()
+_ctx.__enter__()
+_stage_2_transformer = _ctx.__dict__.get('transformer') or \
                        getattr(pipeline.stage_2, '_transformer', None)
-pipeline.stage_2._transformer_ctx = type('ctx', (), {
-    '__enter__': lambda s: _stage_2_transformer,
-    '__exit__': lambda s, *a: None,
-    '__call__': lambda s, *a, **kw: _stage_2_transformer(*a, **kw)
-})()
+pipeline.stage_2._transformer_ctx = lambda: _ctx
 print(f"    Captured stage 2 transformer: {type(_stage_2_transformer)}")
 
-# 2. Load text encoder via _text_encoder_ctx
+# 2. Load text encoder via _text_encoder_ctx() (call first, then enter)
 print("  Loading Gemma text encoder...")
-pipeline.prompt_encoder._text_encoder_ctx.__enter__()
-_text_encoder = pipeline.prompt_encoder._text_encoder_ctx.__dict__.get('text_encoder') or \
+_ctx = pipeline.prompt_encoder._text_encoder_ctx()
+_ctx.__enter__()
+_text_encoder = _ctx.__dict__.get('text_encoder') or \
                 getattr(pipeline.prompt_encoder, '_text_encoder', None)
-# Store as instance attribute and create replacement context
+# Store as instance attribute and create replacement lambda
 pipeline.prompt_encoder._text_encoder = _text_encoder
-pipeline.prompt_encoder._text_encoder_ctx = type('ctx', (), {
-    '__enter__': lambda s: _text_encoder,
-    '__exit__': lambda s, *a: None
-})()
+pipeline.prompt_encoder._text_encoder_ctx = lambda: _ctx
 print(f"    Captured text encoder: {type(_text_encoder)}")
 
-# 3. Load video encoder (from prompt_encoder's video_encoder method)
+# 3. Load video encoder (builder method - returns model directly)
 print("  Loading video encoder...")
 _video_encoder = pipeline.prompt_encoder.video_encoder()
 pipeline.prompt_encoder.video_encoder = lambda: _video_encoder
 print(f"    Captured video encoder: {type(_video_encoder)}")
 
-# 4. Load video decoder via _decoder_builder
+# 4. Load video decoder (builder method)
 print("  Loading video decoder...")
 _video_decoder = pipeline.video_decoder._decoder_builder()
 pipeline.video_decoder._decoder_builder = lambda: _video_decoder
-# Also try direct model attribute if exists
 if hasattr(pipeline.video_decoder, '_decoder'):
     pipeline.video_decoder._decoder = _video_decoder
 print(f"    Captured video decoder: {type(_video_decoder)}")
 
-# 5. Load audio decoder via _decoder_builder
+# 5. Load audio decoder (builder method)
 print("  Loading audio decoder...")
 _audio_decoder = pipeline.audio_decoder._decoder_builder()
 pipeline.audio_decoder._decoder_builder = lambda: _audio_decoder
@@ -247,18 +237,17 @@ if hasattr(pipeline.audio_decoder, '_decoder'):
     pipeline.audio_decoder._decoder = _audio_decoder
 print(f"    Captured audio decoder: {type(_audio_decoder)}")
 
-# 6. Load vocoder (audio decoder has _vocoder_builder)
+# 6. Load vocoder (builder method)
 print("  Loading vocoder...")
 if hasattr(pipeline.audio_decoder, '_vocoder_builder'):
     _vocoder = pipeline.audio_decoder._vocoder_builder()
     pipeline.audio_decoder._vocoder_builder = lambda: _vocoder
     print(f"    Captured vocoder: {type(_vocoder)}")
 
-# 7. Load spatial upsampler via _upsampler_builder
+# 7. Load spatial upsampler (builder method)
 print("  Loading spatial upsampler...")
 _spatial_upsampler = pipeline.upsampler._upsampler_builder()
 pipeline.upsampler._upsampler_builder = lambda: _spatial_upsampler
-# Also try _encoder_builder
 if hasattr(pipeline.upsampler, '_encoder'):
     pipeline.upsampler._encoder = _spatial_upsampler
 print(f"    Captured spatial upsampler: {type(_spatial_upsampler)}")
@@ -270,10 +259,7 @@ if hasattr(pipeline, 'image_conditioner'):
         _ic_encoder = pipeline.image_conditioner.video_encoder()
         pipeline.image_conditioner.video_encoder = lambda: _ic_encoder
 
-# Create references to prevent garbage collection
-# At module level, variables are already global - no 'global' keyword needed
 print("  Models captured and preserved for ZeroGPU tensor packing...")
-
 print("All models preloaded for ZeroGPU tensor packing!")
 print("=" * 80)