Fix issue with layer offloading with MinMax H3

extensions_built_in/diffusion_models/minimax_h3/src/transformer.py

@@ -87,8 +87,11 @@ class MiniMaxH3Rope(nn.Module):
 
     def forward(self, position_ids: torch.Tensor):
         """position_ids (B, S, 3) -> cos, sin each (B, S, 96), float32."""
-        position_ids = position_ids.to(device=self.inv_freq.device, dtype=torch.float32)
-        freqs = position_ids.unsqueeze(-1) * self.inv_freq.view(1, 1, 1, -1)
+        # compute on the input's device: the frequency buffer may be left
+        # CPU-resident by layer offloading / low_vram loads
+        position_ids = position_ids.to(dtype=torch.float32)
+        inv_freq = self.inv_freq.to(position_ids.device)
+        freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, 1, -1)
         # (B, S, 3, 16) -> (B, S, 48) in (t, h, w) axis order -> duplicate to 96
         freqs = freqs.flatten(2, 3)
         freqs = torch.cat([freqs, freqs], dim=-1)
@@ -403,10 +406,13 @@ class MiniMaxH3Transformer(nn.Module):
             return self.time_embedder(t)
         table = self.adaln_t_table.float()
         pos = t.clamp(0.0, 1.0) * (table.shape[0] - 1)
+        # the table may stay CPU-resident under offloading while timesteps
+        # arrive on cuda: interpolate on the table's device, return on t's
+        pos = pos.to(table.device)
         lo = pos.floor().long()
         hi = (lo + 1).clamp(max=table.shape[0] - 1)
         frac = (pos - lo.float()).unsqueeze(1)
-        return table[lo] * (1.0 - frac) + table[hi] * frac
+        return (table[lo] * (1.0 - frac) + table[hi] * frac).to(t.device)
 
     def forward(
         self,

toolkit/util/comfy_quant_import.py

@@ -62,11 +62,13 @@ class Int8Embedding(torch.nn.Module):
         return (self.qweight.float() * scales.unsqueeze(1)).to(self.output_dtype)
 
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        flat = input_ids.reshape(-1)
+        # the table may stay CPU-resident under text-encoder offloading: run
+        # the (tiny) lookup on the table's device, return on the caller's
+        flat = input_ids.reshape(-1).to(self.qweight.device)
         rows = self.qweight.index_select(0, flat).float()
         scales = self.scales.view(torch.float32).index_select(0, flat)
-        out = rows * scales.unsqueeze(1)
-        return out.to(self.output_dtype).reshape(*input_ids.shape, self.embedding_dim)
+        out = (rows * scales.unsqueeze(1)).to(self.output_dtype)
+        return out.to(input_ids.device).reshape(*input_ids.shape, self.embedding_dim)
 
 
 def _to_ostris(module: torch.nn.Linear, quantizer, orig_dtype: torch.dtype) -> OstrisLinear: