github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix issue with layer offloading with MinMax H3

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-03 11:32:36 -06:00
parent 9d614a51fb
commit bf739ff966
2 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
toolkit/util/comfy_quant_import.py
View File

@@ -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: