Build on CPU, move to cuda once (ZeroGPU)

zoom.py

@@ -5,6 +5,7 @@ of them to the model: no disk round-trip between recursions, the zoom window can
 anywhere instead of only the centre, and the loop is a generator so the UI can show each
 level the moment it lands. Everything on one device, which is what ZeroGPU gives us.
 """
+import gc
 import os
 import sys
 
@@ -60,7 +61,10 @@ def _build_vlm(lora_path):
     model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         VLM, torch_dtype=torch.bfloat16, attn_implementation="sdpa"
     )
-    model = PeftModel.from_pretrained(model, lora_path).merge_and_unload()
+    # torch_device="cpu" is load-bearing. Left to itself peft calls infer_device(), which reads
+    # the cuda flag ZeroGPU patches to True, then hands "cuda" to safetensors. That builds
+    # tensors straight on a device that is not attached yet: RuntimeError, no CUDA GPUs.
+    model = PeftModel.from_pretrained(model, lora_path, torch_device="cpu").merge_and_unload()
     return model.eval().to("cuda"), AutoProcessor.from_pretrained(VLM), process_vision_info
 
 
@@ -68,13 +72,10 @@ def _build_sr(coz_ckpt, merged_transformer):
     from safetensors.torch import load_file
     from osediff_sd3 import OSEDiff_SD3_TEST, SD3Euler
 
-    sr = SD3Euler()
-    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3):
-        m.to("cuda")
-    sr.transformer.to("cuda", dtype=torch.float32)
-    sr.vae.to("cuda", dtype=torch.float32)
-    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae):
-        m.requires_grad_(False)
+    # Assemble entirely on CPU, then move once. Two places here materialise weights straight
+    # onto the model's own device (inject_lora's torch.load map_location, and safetensors), and
+    # on ZeroGPU that device reads as cuda while no GPU is attached yet.
+    sr = SD3Euler(device="cpu")
     # Construct first, load second. OSEDiff_SD3_TEST swaps every targeted Linear for a
     # LoraInjectedLinear, which renames the keys. Our merged checkpoint was saved after that
     # swap, so loading it earlier matches nothing and silently leaves the base transformer.
@@ -86,6 +87,16 @@ def _build_sr(coz_ckpt, merged_transformer):
         raise RuntimeError(f"merged transformer did not match: {len(unexpected)} of {len(sd)} "
                            f"keys unexpected. LoRA injection order or checkpoint is wrong.")
     print(f"#### merged transformer loaded (missing {len(missing)} unexpected {len(unexpected)})")
+    del sd
+    gc.collect()
+
+    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3):
+        m.to("cuda")
+    sr.transformer.to("cuda", dtype=torch.float32)
+    sr.vae.to("cuda", dtype=torch.float32)
+    for m in (sr.text_enc_1, sr.text_enc_2, sr.text_enc_3, sr.transformer, sr.vae):
+        m.requires_grad_(False)
+    sr.device = "cuda"   # encode_prompt and set_timesteps read this, not the module devices
     return test