Update app.py

app.py

@@ -178,83 +178,86 @@ print("=" * 80)
 # =============================================================================
 # ZeroGPU Tensor Preloading
 # =============================================================================
-print("Preloading all models for ZeroGPU tensor packing...")
+# NOTE: At Space startup, no GPU is available (ZeroGPU assigns it at runtime).
+# We can only preload components that don't require CUDA.
+# The transformer (and other GPU-heavy components) will load during generation
+# when ZeroGPU provides the GPU. ZeroGPU should capture them then.
+
+print("Preloading non-CUDA components for ZeroGPU tensor packing...")
 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 (require streaming_prefetch_count)
-# We need to call these methods, capture the results, and preserve them
-
-# 1. Load transformer via _transformer_ctx(streaming_prefetch_count=None)
-print("  Loading stage 1 transformer...")
-_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)
-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(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 streaming_prefetch_count=None: _ctx
-print(f"    Captured stage 2 transformer: {type(_stage_2_transformer)}")
-
-# 2. Load text encoder via _text_encoder_ctx(streaming_prefetch_count=None)
-print("  Loading Gemma text encoder...")
-_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)
-pipeline.prompt_encoder._text_encoder = _text_encoder
-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)
+# 1. Try loading video encoder (may work without GPU if just file loading)
 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 (builder method)
+try:
+    _video_encoder = pipeline.prompt_encoder.video_encoder()
+    pipeline.prompt_encoder.video_encoder = lambda: _video_encoder
+    print(f"    Loaded video encoder: {type(_video_encoder)}")
+except Exception as e:
+    print(f"    Video encoder preload skipped: {e}")
+
+# 2. Try loading video decoder (VAE - may work without GPU)
 print("  Loading video decoder...")
-_video_decoder = pipeline.video_decoder._decoder_builder()
-pipeline.video_decoder._decoder_builder = lambda: _video_decoder
-if hasattr(pipeline.video_decoder, '_decoder'):
-    pipeline.video_decoder._decoder = _video_decoder
-print(f"    Captured video decoder: {type(_video_decoder)}")
-
-# 5. Load audio decoder (builder method)
+try:
+    _video_decoder = pipeline.video_decoder._decoder_builder()
+    pipeline.video_decoder._decoder_builder = lambda: _video_decoder
+    if hasattr(pipeline.video_decoder, '_decoder'):
+        pipeline.video_decoder._decoder = _video_decoder
+    print(f"    Loaded video decoder: {type(_video_decoder)}")
+except Exception as e:
+    print(f"    Video decoder preload skipped: {e}")
+
+# 3. Try loading audio decoder (VAE - may work without GPU)
 print("  Loading audio decoder...")
-_audio_decoder = pipeline.audio_decoder._decoder_builder()
-pipeline.audio_decoder._decoder_builder = lambda: _audio_decoder
-if hasattr(pipeline.audio_decoder, '_decoder'):
-    pipeline.audio_decoder._decoder = _audio_decoder
-print(f"    Captured audio decoder: {type(_audio_decoder)}")
-
-# 6. Load vocoder (builder method)
+try:
+    _audio_decoder = pipeline.audio_decoder._decoder_builder()
+    pipeline.audio_decoder._decoder_builder = lambda: _audio_decoder
+    if hasattr(pipeline.audio_decoder, '_decoder'):
+        pipeline.audio_decoder._decoder = _audio_decoder
+    print(f"    Loaded audio decoder: {type(_audio_decoder)}")
+except Exception as e:
+    print(f"    Audio decoder preload skipped: {e}")
+
+# 4. Try loading vocoder
 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 (builder method)
+try:
+    if hasattr(pipeline.audio_decoder, '_vocoder_builder'):
+        _vocoder = pipeline.audio_decoder._vocoder_builder()
+        pipeline.audio_decoder._vocoder_builder = lambda: _vocoder
+        print(f"    Loaded vocoder: {type(_vocoder)}")
+except Exception as e:
+    print(f"    Vocoder preload skipped: {e}")
+
+# 5. Try loading spatial upsampler
 print("  Loading spatial upsampler...")
-_spatial_upsampler = pipeline.upsampler._upsampler_builder()
-pipeline.upsampler._upsampler_builder = lambda: _spatial_upsampler
-if hasattr(pipeline.upsampler, '_encoder'):
-    pipeline.upsampler._encoder = _spatial_upsampler
-print(f"    Captured spatial upsampler: {type(_spatial_upsampler)}")
-
-# 8. Load image conditioner
+try:
+    _spatial_upsampler = pipeline.upsampler._upsampler_builder()
+    pipeline.upsampler._upsampler_builder = lambda: _spatial_upsampler
+    if hasattr(pipeline.upsampler, '_encoder'):
+        pipeline.upsampler._encoder = _spatial_upsampler
+    print(f"    Loaded spatial upsampler: {type(_spatial_upsampler)}")
+except Exception as e:
+    print(f"    Spatial upsampler preload skipped: {e}")
+
+# 6. Load image conditioner
 print("  Loading image conditioner...")
-if hasattr(pipeline, 'image_conditioner'):
-    if hasattr(pipeline.image_conditioner, 'video_encoder'):
-        _ic_encoder = pipeline.image
+try:
+    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(f"    Loaded image conditioner encoder")
+except Exception as e:
+    print(f"    Image conditioner preload skipped: {e}")
+
+# 7. NOTE: Transformer loading is intentionally skipped here
+# The transformer requires CUDA (LoRA fusion uses triton kernels)
+# It will load during generate_video() when ZeroGPU provides a GPU
+# ZeroGPU should capture it then
+print("  Transformer: Will load during generation (requires GPU)")
+print("  Text encoder: Will load during generation (requires GPU)")
+
+print("Non-CUDA components preloaded!")
+print("=" * 80)
 
 # =============================================================================
 # Helper Functions