Upload folder using huggingface_hub

README.md

@@ -33,7 +33,7 @@ from transformers import AutoModel
 import torch
 
 model = AutoModel.from_pretrained(
-    "QuarkML/QaDiT",          # or a local export folder
+    "USER/qadit",          # or a local export folder
     trust_remote_code=True,
 ).to("cuda" if torch.cuda.is_available() else "cpu")
 

modelling_qadit.py

@@ -353,19 +353,46 @@ class DiffusionScheduler:
         self.T = num_train_steps
         self.ln_mean = logit_normal_mean
         self.ln_std = logit_normal_std
+        self.schedule = schedule
         if schedule != "cosine":
             raise ValueError(f"unknown schedule: {schedule}")
+        # May be created on the meta device under HF's init_empty_weights();
+        # materialize_real() / to() rebuilds a real CPU/CUDA table before use.
+        self.alpha_bar = self._build_alpha_bar(self.T)
+
+    @staticmethod
+    def _build_alpha_bar(num_train_steps: int, device=None) -> torch.Tensor:
+        device = torch.device(device) if device is not None else torch.device("cpu")
+        # Force a concrete device — never allocate on "meta".
+        if device.type == "meta":
+            device = torch.device("cpu")
         s = 0.008
-        steps = torch.arange(self.T + 1, dtype=torch.float64)
-        f = torch.cos((steps / self.T + s) / (1 + s) * math.pi / 2) ** 2
+        steps = torch.arange(
+            num_train_steps + 1, dtype=torch.float64, device=device
+        )
+        f = torch.cos((steps / num_train_steps + s) / (1 + s) * math.pi / 2) ** 2
         abar = (f / f[0]).clamp(1e-5, 1.0)
-        self.alpha_bar = abar[1:].float()
+        return abar[1:].float()
+
+    def _is_meta(self) -> bool:
+        t = self.alpha_bar
+        return bool(getattr(t, "is_meta", False) or t.device.type == "meta")
+
+    def materialize_real(self, device=None) -> "DiffusionScheduler":
+        """Rebuild alpha_bar if it was left on the meta device by from_pretrained."""
+        target = torch.device(device) if device is not None else torch.device("cpu")
+        if target.type == "meta":
+            target = torch.device("cpu")
+        if self._is_meta() or self.alpha_bar.device != target:
+            self.alpha_bar = self._build_alpha_bar(self.T, device="cpu").to(target)
+        return self
 
     def to(self, device) -> "DiffusionScheduler":
-        self.alpha_bar = self.alpha_bar.to(device)
-        return self
+        return self.materialize_real(device)
 
     def _gather(self, t: torch.Tensor):
+        if self._is_meta():
+            self.materialize_real(t.device)
         abar = self.alpha_bar.to(t.device)[t]
         return abar.sqrt().view(-1, 1, 1, 1), (1 - abar).sqrt().view(-1, 1, 1, 1)
 
@@ -390,6 +417,7 @@ class DiffusionScheduler:
         device: Union[str, torch.device] = "cpu",
         generator: Optional[torch.Generator] = None,
     ) -> torch.Tensor:
+        self.materialize_real(device)
         B = shape[0]
         z = torch.randn(shape, device=device, generator=generator)
         times = torch.linspace(self.T - 1, 0, num_steps, device=device).long()
@@ -411,8 +439,8 @@ class DiffusionScheduler:
                 break
 
             t_next = times[i + 1].expand(B)
-            abar_next = self.alpha_bar.to(device)[t_next].view(-1, 1, 1, 1)
-            abar_now = self.alpha_bar.to(device)[t].view(-1, 1, 1, 1)
+            abar_next = self.alpha_bar[t_next].view(-1, 1, 1, 1)
+            abar_now = self.alpha_bar[t].view(-1, 1, 1, 1)
             sigma = eta * torch.sqrt(
                 (1 - abar_next) / (1 - abar_now) * (1 - abar_now / abar_next)
             )