SDNQ add Flux2Transformer2DModel keys

This commit is contained in:
Disty0
2025-11-25 19:21:59 +03:00
parent 81e44ca90f
commit aeb71d172e
+7 -2
View File
@@ -82,6 +82,9 @@ else:
int_mm_func = torch._int_mm
def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
return torch.mm(x,y, out_dtype=torch.float32)
fp_mm_func = None
if os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
try:
@@ -91,8 +94,6 @@ if os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}
fp_mm_func = None
if fp_mm_func is None:
def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
return torch.mm(x,y, out_dtype=torch.float32)
fp_mm_func = fp_mm_torch
@@ -139,6 +140,10 @@ module_skip_keys_dict = {
["single_transformer_blocks.0.norm.linear.weight", "time_text_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{}
],
"Flux2Transformer2DModel": [
["double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation", "time_guidance_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{}
],
"ChromaTransformer2DModel": [
["distilled_guidance_layer", "time_text_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{}