273 lines
10 KiB
Python
273 lines
10 KiB
Python
import torch
|
|
import os
|
|
from diffusers import DiffusionPipeline, DDIMScheduler, Transformer2DModel, AutoencoderKL
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
|
|
from safetensors.torch import load_file
|
|
from PIL import Image
|
|
import numpy as np
|
|
import quanto
|
|
|
|
# --- Modular Loaders ---
|
|
|
|
class LoadQuantoZImageTransformer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
|
|
"transformer_path": ("STRING", {"default": "models/quantized_models/zimage_turbo_transformer_qdit.safetensors"}),
|
|
"dtype": (["bfloat16", "float16"], {"default": "bfloat16"}),
|
|
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ZIMAGE_TRANSFORMER",)
|
|
RETURN_NAMES = ("transformer",)
|
|
FUNCTION = "load_transformer"
|
|
CATEGORY = "Z-Image (Modular)"
|
|
|
|
def load_transformer(self, model_id, transformer_path, dtype, device):
|
|
print(f"Loading Quantized Transformer from {transformer_path}...")
|
|
torch_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
|
|
|
if device == "auto":
|
|
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
else:
|
|
dev = torch.device(device)
|
|
|
|
# Try to import ZImageTransformer2DModel
|
|
try:
|
|
from diffusers import ZImageTransformer2DModel
|
|
TransformerClass = ZImageTransformer2DModel
|
|
except ImportError:
|
|
print("ZImageTransformer2DModel not found. Using Transformer2DModel.")
|
|
TransformerClass = Transformer2DModel
|
|
|
|
try:
|
|
config = TransformerClass.load_config(model_id, subfolder="transformer")
|
|
with torch.device("meta"):
|
|
transformer = TransformerClass.from_config(config)
|
|
except Exception as e:
|
|
print(f"Failed to load transformer config: {e}")
|
|
raise e
|
|
|
|
transformer.to_empty(device="cpu")
|
|
|
|
print("Initializing Quantized Structure (quanto)...")
|
|
modules_to_quantize = []
|
|
for name, module in transformer.named_modules():
|
|
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
|
|
modules_to_quantize.append(module)
|
|
|
|
for module in modules_to_quantize:
|
|
quanto.quantize(module, weights=quanto.qint4, activations=quanto.qint8)
|
|
|
|
print("Loading Weights (Manual Assignment)...")
|
|
state_dict = load_file(transformer_path)
|
|
|
|
# Manually assign weights to bypass strict checks in quanto's load_state_dict hook
|
|
# This is necessary because we filtered out metadata keys during saving
|
|
model_dict = transformer.state_dict()
|
|
for name, param in transformer.named_parameters():
|
|
if name in state_dict:
|
|
# We assume the shapes match.
|
|
# If param is a QTensor, we might need to handle it carefully,
|
|
# but usually assigning to .data works if the tensor is compatible.
|
|
try:
|
|
param.data = state_dict[name].to(param.device)
|
|
except Exception as e:
|
|
print(f"Failed to assign {name}: {e}")
|
|
elif name.endswith("_scale") and name in state_dict:
|
|
# Handle scales if they are separate parameters/buffers
|
|
pass
|
|
|
|
# Also handle buffers (like scales if they are buffers)
|
|
for name, buf in transformer.named_buffers():
|
|
if name in state_dict:
|
|
try:
|
|
buf.data = state_dict[name].to(buf.device)
|
|
except Exception as e:
|
|
print(f"Failed to assign buffer {name}: {e}")
|
|
|
|
del state_dict
|
|
|
|
quanto.freeze(transformer)
|
|
transformer.to(dev)
|
|
|
|
return (transformer,)
|
|
|
|
class LoadQuantoZImageCLIP:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
|
|
"text_encoder_path": ("STRING", {"default": "models/quantized_models/qwen_text_encoder_qdit.safetensors"}),
|
|
"dtype": (["bfloat16", "float16"], {"default": "bfloat16"}),
|
|
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ZIMAGE_CLIP",)
|
|
RETURN_NAMES = ("text_encoder",)
|
|
FUNCTION = "load_clip"
|
|
CATEGORY = "Z-Image (Modular)"
|
|
|
|
def load_clip(self, model_id, text_encoder_path, dtype, device):
|
|
print(f"Loading Quantized Text Encoder from {text_encoder_path}...")
|
|
torch_dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float16
|
|
|
|
if device == "auto":
|
|
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
else:
|
|
dev = torch.device(device)
|
|
|
|
print("Loading Config (Qwen3)...")
|
|
try:
|
|
config = AutoConfig.from_pretrained("Qwen/Qwen3-4B")
|
|
except Exception:
|
|
print("Warning: Could not load Qwen/Qwen3-4B config directly. Trying from model_id...")
|
|
config = AutoConfig.from_pretrained(model_id, subfolder="text_encoder")
|
|
|
|
with torch.device("meta"):
|
|
text_encoder = AutoModelForCausalLM.from_config(config)
|
|
|
|
text_encoder.to_empty(device="cpu")
|
|
|
|
print("Initializing Quantized Structure (quanto)...")
|
|
modules_to_quantize = []
|
|
for name, module in text_encoder.named_modules():
|
|
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
|
|
modules_to_quantize.append(module)
|
|
|
|
for module in modules_to_quantize:
|
|
quanto.quantize(module, weights=quanto.qint4, activations=quanto.qint8)
|
|
|
|
print("Loading Weights (Manual Assignment)...")
|
|
state_dict = load_file(text_encoder_path)
|
|
|
|
# Manually assign weights to bypass strict checks
|
|
for name, param in text_encoder.named_parameters():
|
|
if name in state_dict:
|
|
try:
|
|
param.data = state_dict[name].to(param.device)
|
|
except Exception as e:
|
|
print(f"Failed to assign {name}: {e}")
|
|
|
|
for name, buf in text_encoder.named_buffers():
|
|
if name in state_dict:
|
|
try:
|
|
buf.data = state_dict[name].to(buf.device)
|
|
except Exception as e:
|
|
print(f"Failed to assign buffer {name}: {e}")
|
|
|
|
del state_dict
|
|
|
|
quanto.freeze(text_encoder)
|
|
text_encoder.to(dev, dtype=torch_dtype)
|
|
|
|
return (text_encoder,)
|
|
|
|
class LoadQuantoZImageVAE:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
|
|
"dtype": (["bfloat16", "float16", "float32"], {"default": "bfloat16"}),
|
|
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ZIMAGE_VAE",)
|
|
RETURN_NAMES = ("vae",)
|
|
FUNCTION = "load_vae"
|
|
CATEGORY = "Z-Image (Modular)"
|
|
|
|
def load_vae(self, model_id, dtype, device):
|
|
print(f"Loading VAE from {model_id}...")
|
|
torch_dtype = torch.bfloat16 if dtype == "bfloat16" else (torch.float16 if dtype == "float16" else torch.float32)
|
|
|
|
if device == "auto":
|
|
dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
else:
|
|
dev = torch.device(device)
|
|
|
|
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=torch_dtype)
|
|
vae.to(dev)
|
|
|
|
return (vae,)
|
|
|
|
class ZImageQuantoSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"transformer": ("ZIMAGE_TRANSFORMER",),
|
|
"text_encoder": ("ZIMAGE_CLIP",),
|
|
"vae": ("ZIMAGE_VAE",),
|
|
"prompt": ("STRING", {"default": "a beautiful landscape"}),
|
|
"negative_prompt": ("STRING", {"default": ""}),
|
|
"height": ("INT", {"default": 1024, "min": 256, "max": 2048}),
|
|
"width": ("INT", {"default": 1024, "min": 256, "max": 2048}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 100}),
|
|
"guidance_scale": ("FLOAT", {"default": 7.5}),
|
|
"seed": ("INT", {"default": 42}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Z-Image (Modular)"
|
|
|
|
def generate(self, transformer, text_encoder, vae, prompt, negative_prompt, height, width, steps, guidance_scale, seed):
|
|
torch.manual_seed(seed)
|
|
device = transformer.device
|
|
|
|
# Assemble a temporary pipeline for generation
|
|
# We need the tokenizer too. Ideally this should be passed or loaded here.
|
|
# For simplicity, we load it here (lightweight).
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
|
|
except Exception:
|
|
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-7B", local_files_only=False)
|
|
|
|
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
|
|
|
|
# We construct the pipeline
|
|
# Note: We can't easily pass this to standard KSampler because it's a diffusers pipeline.
|
|
# So we run the loop here.
|
|
|
|
pipe = DiffusionPipeline.from_pretrained(
|
|
"Tongyi-MAI/Z-Image-Turbo", # Dummy ID to get structure if needed, or just use components
|
|
text_encoder=text_encoder,
|
|
transformer=transformer,
|
|
vae=vae,
|
|
tokenizer=tokenizer,
|
|
scheduler=scheduler,
|
|
torch_dtype=transformer.dtype,
|
|
device_map=None
|
|
)
|
|
pipe.to(device)
|
|
|
|
print("Generating image...")
|
|
try:
|
|
image = pipe(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=steps,
|
|
guidance_scale=guidance_scale,
|
|
output_type="pil"
|
|
).images[0]
|
|
except Exception as e:
|
|
print(f"Pipeline call failed: {e}.")
|
|
raise e
|
|
|
|
image_np = np.array(image)
|
|
image_tensor = torch.from_numpy(image_np).float() / 255.0
|
|
image_tensor = image_tensor.unsqueeze(0)
|
|
|
|
return (image_tensor,)
|