Update app.py

app.py

@@ -183,36 +183,34 @@ print("This may take a few minutes...")
 
 # TI2VidTwoStagesHQPipeline uses:
 # - Builder methods that return models directly when called
-# - Context methods that return context managers when called
+# - Context methods that return context managers when called (require streaming_prefetch_count)
 # We need to call these methods, capture the results, and preserve them
 
-# 1. Load transformer via _transformer_ctx() (call first, then enter)
+# 1. Load transformer via _transformer_ctx(streaming_prefetch_count=None)
 print("  Loading stage 1 transformer...")
-_ctx = pipeline.stage_1._transformer_ctx()  # Get context manager
-_ctx.__enter__()  # Enter context
+_ctx = pipeline.stage_1._transformer_ctx(streaming_prefetch_count=None)
+_ctx.__enter__()
 _stage_1_transformer = _ctx.__dict__.get('transformer') or \
                        getattr(pipeline.stage_1, '_transformer', None)
-# Replace _transformer_ctx with a lambda that returns cached model
-pipeline.stage_1._transformer_ctx = lambda: _ctx
+pipeline.stage_1._transformer_ctx = lambda streaming_prefetch_count=None: _ctx
 print(f"    Captured stage 1 transformer: {type(_stage_1_transformer)}")
 
 print("  Loading stage 2 transformer...")
-_ctx = pipeline.stage_2._transformer_ctx()
+_ctx = pipeline.stage_2._transformer_ctx(streaming_prefetch_count=None)
 _ctx.__enter__()
 _stage_2_transformer = _ctx.__dict__.get('transformer') or \
                        getattr(pipeline.stage_2, '_transformer', None)
-pipeline.stage_2._transformer_ctx = lambda: _ctx
+pipeline.stage_2._transformer_ctx = lambda streaming_prefetch_count=None: _ctx
 print(f"    Captured stage 2 transformer: {type(_stage_2_transformer)}")
 
-# 2. Load text encoder via _text_encoder_ctx() (call first, then enter)
+# 2. Load text encoder via _text_encoder_ctx(streaming_prefetch_count=None)
 print("  Loading Gemma text encoder...")
-_ctx = pipeline.prompt_encoder._text_encoder_ctx()
+_ctx = pipeline.prompt_encoder._text_encoder_ctx(streaming_prefetch_count=None)
 _ctx.__enter__()
 _text_encoder = _ctx.__dict__.get('text_encoder') or \
                 getattr(pipeline.prompt_encoder, '_text_encoder', None)
-# Store as instance attribute and create replacement lambda
 pipeline.prompt_encoder._text_encoder = _text_encoder
-pipeline.prompt_encoder._text_encoder_ctx = lambda: _ctx
+pipeline.prompt_encoder._text_encoder_ctx = lambda streaming_prefetch_count=None: _ctx
 print(f"    Captured text encoder: {type(_text_encoder)}")
 
 # 3. Load video encoder (builder method - returns model directly)
@@ -256,12 +254,7 @@ print(f"    Captured spatial upsampler: {type(_spatial_upsampler)}")
 print("  Loading image conditioner...")
 if hasattr(pipeline, 'image_conditioner'):
     if hasattr(pipeline.image_conditioner, 'video_encoder'):
-        _ic_encoder = pipeline.image_conditioner.video_encoder()
-        pipeline.image_conditioner.video_encoder = lambda: _ic_encoder
-
-print("  Models captured and preserved for ZeroGPU tensor packing...")
-print("All models preloaded for ZeroGPU tensor packing!")
-print("=" * 80)
+        _ic_encoder = pipeline.image
 
 # =============================================================================
 # Helper Functions