Never .to() a torchao-quantised param (ZeroGPU empty_like on Float8Tensor)

xvideo/models/loader.py

@@ -102,6 +102,20 @@ def _is_fp8_target(module: nn.Module, fqn: str) -> bool:
     return any(sub in fqn for sub in _FP8_TARGET_SUBSTRINGS)
 
 
+def _move_unquantised(model: nn.Module, device: str) -> None:
+    """Move every top-level child except ``double_blocks`` (already on device)."""
+    for name, child in model.named_children():
+        if name == "double_blocks":
+            continue
+        child.to(device)
+    for name, param in list(model.named_parameters(recurse=False)):
+        if param.device.type != torch.device(device).type:
+            setattr(model, name, nn.Parameter(param.data.to(device), requires_grad=False))
+    for name, buf in list(model.named_buffers(recurse=False)):
+        if buf is not None and buf.device.type != torch.device(device).type:
+            model.register_buffer(name, buf.to(device))
+
+
 def build_dit(cfg: ExpConfig, dit_path: str, device: str = "cuda") -> nn.Module:
     dtype = PRECISION_TO_TYPE[cfg.dit_precision]
     params = dict(cfg.dit_arch_config.get("params", {}))
@@ -153,15 +167,18 @@ def build_dit(cfg: ExpConfig, dit_path: str, device: str = "cuda") -> nn.Module:
                 _log(f"  block {idx + 1}/{n_blocks}")
         _log(f"FP8 quantisation done in {time.time() - t0:.0f}s")
 
-    # Anything still mmap-backed (patch embed, modulation, norms, proj_out) is
-    # cloned so the 32.5 GB file can be released, then moved to the device.
-    for name, param in model.named_parameters():
-        if param.device.type == "cpu":
-            param.data = param.data.clone()
+        # The blocks are already on the device.  Everything still mmap-backed
+        # (patch embed, timestep/text projections, final layer) is moved here.
+        # A quantised parameter must NEVER be moved again: `Tensor.to` on a
+        # torchao `Float8Tensor` hits an unimplemented `aten.empty_like` inside
+        # the ZeroGPU tensor-tracking hook.
+        _move_unquantised(model, device)
+    else:
+        model.to(device)
+
     del state_dict
     gc.collect()
 
-    model.to(device)
     if getattr(model.config, "causal", False):
         model.config.use_inference_kv_cache = True
     return model