diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py b/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py index 9fc0bdb4..53501065 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/transformer.py +++ b/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, diff --git a/toolkit/util/comfy_quant_import.py b/toolkit/util/comfy_quant_import.py index 85f31994..833d6b2d 100644 --- a/toolkit/util/comfy_quant_import.py +++ b/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: