fix: load weights on CPU then move modules to cuda (ZeroGPU can't intercept safetensors device=cuda)

app.py

@@ -163,41 +163,62 @@ from ltx_distillation.utils import (  # noqa: E402
 from ltx_pipelines.utils.model_ledger import ModelLedger  # noqa: E402
 
 DEVICE = torch.device("cuda")
-
-logger.info("Building model ledger …")
+CPU = torch.device("cpu")
+
+# The ledger is built on CPU and every module is moved to "cuda" afterwards.
+# That is the authors' own pattern — `create_ltx2_wrapper` / `create_vae_wrappers`
+# / `create_text_encoder_wrapper` all construct a `device=torch.device("cpu")`
+# ledger and then `.to(device)`, because `ModelLedger._target_device()` is handed
+# straight to `safetensors.safe_open(..., device=...)`. Loading weights inside
+# safetensors bypasses ZeroGPU's `torch.cuda` interception, so asking it for
+# "cuda" in the main process dies with "No CUDA GPUs are available". Moving the
+# assembled `nn.Module`s instead goes through the hijacked `Module.to`, which is
+# exactly what ZeroGPU wants: weights get packed and streamed into VRAM on the
+# first `@spaces.GPU` call.
+logger.info("Building model ledger (weights load on CPU, then move to cuda) …")
 _ledger = ModelLedger(
     dtype=DTYPE,
-    device=DEVICE,
+    device=CPU,
     checkpoint_path=CHECKPOINT_PATH,
     gemma_root_path=GEMMA_ROOT,
     quantization=QuantizationPolicy.fp8_cast(),  # transformer only
 )
 
 logger.info("Loading audio-video DiT (fp8 cast) …")
-_transformer = _ledger.transformer()
+# NOTE: no `dtype=` here, unlike the authors' unquantized `create_ltx2_wrapper`.
+# `QuantizationPolicy.fp8_cast()` has already stored the transformer-block linear
+# weights as float8_e4m3fn; passing dtype=bfloat16 would upcast them straight back
+# and undo the memory saving that lets this model fit at all.
+_transformer = _ledger.transformer().to(device=DEVICE)
+
 logger.info("Loading video VAE …")
 _video_vae = VideoVAEWrapper(
-    encoder=_ledger.video_encoder(),
-    decoder=_ledger.video_decoder(),
+    encoder=_ledger.video_encoder().to(device=DEVICE, dtype=DTYPE),
+    decoder=_ledger.video_decoder().to(device=DEVICE, dtype=DTYPE),
     device=DEVICE,
     dtype=DTYPE,
 ).eval()
+
 logger.info("Loading audio VAE + vocoder …")
 _audio_vae = AudioVAEWrapper(
-    encoder=_ledger.audio_encoder(),
-    decoder=_ledger.audio_decoder(),
-    vocoder=_ledger.vocoder(),
+    # float32 for the audio encoder, matching `create_vae_wrappers`.
+    encoder=_ledger.audio_encoder().to(device=DEVICE, dtype=torch.float32),
+    decoder=_ledger.audio_decoder().to(device=DEVICE, dtype=DTYPE),
+    vocoder=_ledger.vocoder().to(device=DEVICE, dtype=DTYPE),
     device=DEVICE,
     dtype=DTYPE,
 ).eval()
 AUDIO_SR = _audio_vae.get_output_sample_rate() or 24000
+
 logger.info("Loading Gemma-3-12B text encoder + embeddings processor …")
 _text_encoder = GemmaTextEncoderWrapper(
-    text_encoder=_ledger.text_encoder(),
-    embeddings_processor=_ledger.gemma_embeddings_processor(),
+    text_encoder=_ledger.text_encoder().to(device=DEVICE, dtype=DTYPE),
+    embeddings_processor=_ledger.gemma_embeddings_processor().to(device=DEVICE, dtype=DTYPE),
     device=DEVICE,
     dtype=DTYPE,
 ).eval()
+
+gc.collect()
 logger.info("All models resident. audio_sample_rate=%d", AUDIO_SR)
 
 _SIGMAS_CPU = torch.tensor(DENOISING_SIGMAS, dtype=torch.float32)