74 lines
3.2 KiB
Python
74 lines
3.2 KiB
Python
|
|
import torch
|
|
|
|
from diffusers import DiffusionPipeline
|
|
from safetensors.torch import load_file
|
|
import os
|
|
|
|
class LoadZImageTurboQDiTCalibratedCompile:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
|
|
"transformer_path": ("STRING", {"default": "quantized_models/zimage_turbo_transformer_qdit_calibrated.safetensors"}),
|
|
"text_encoder_path": ("STRING", {"default": "quantized_models/qwen_text_encoder_qdit_calibrated.safetensors"}),
|
|
"dtype": (["bfloat16", "float16"], {"default": "bfloat16"}),
|
|
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
|
"enable_compile": ("BOOL", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ZIMAGE_PIPELINE",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "Z-Image (Turbo)"
|
|
|
|
def load(self, model_id, transformer_path, text_encoder_path, dtype, device, enable_compile):
|
|
torch_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
|
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype, trust_remote_code=True)
|
|
|
|
dev = torch.device("cuda" if (device == "auto" and torch.cuda.is_available()) or device == "cuda" else "cpu")
|
|
pipe.to(dev)
|
|
|
|
print("Loading calibrated Q-DiT transformer from .safetensors...")
|
|
state_dict_transformer = load_file(transformer_path)
|
|
activation_scales = None
|
|
if "__activation_scales__" in state_dict_transformer:
|
|
activation_scales = state_dict_transformer.pop("__activation_scales__")
|
|
pipe.transformer.load_state_dict(state_dict_transformer)
|
|
|
|
if hasattr(pipe, "text_encoder") and os.path.exists(text_encoder_path):
|
|
print("Loading calibrated Q-DiT text encoder from .safetensors...")
|
|
state_dict_text = load_file(text_encoder_path)
|
|
if "__activation_scales__" in state_dict_text:
|
|
state_dict_text.pop("__activation_scales__")
|
|
pipe.text_encoder.load_state_dict(state_dict_text)
|
|
|
|
# Wrap transformer forward to apply activation scaling
|
|
if activation_scales is not None:
|
|
scales_list = activation_scales.tolist()
|
|
original_forward = pipe.transformer.forward
|
|
|
|
def scaled_forward(*args, **kwargs):
|
|
output = original_forward(*args, **kwargs)
|
|
if isinstance(output, torch.Tensor):
|
|
scale_factor = max(scales_list) if scales_list else 1.0
|
|
output = output / scale_factor
|
|
return output
|
|
|
|
pipe.transformer.forward = scaled_forward
|
|
print("Activation scaling applied during inference.")
|
|
|
|
# Optional torch.compile for DiT transformer
|
|
if enable_compile:
|
|
try:
|
|
print("Compiling transformer with torch.compile (mode=reduce-overhead)...")
|
|
pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=False)
|
|
print("Compilation successful.")
|
|
except Exception as e:
|
|
print(f"torch.compile failed: {e}. Continuing without compilation.")
|
|
|
|
return (pipe,)
|
|
|
|
|