327 lines
13 KiB
Python
327 lines
13 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import subprocess
|
|
import sys
|
|
from diffusers import DiffusionPipeline
|
|
|
|
# Try to import sdnq, install with workaround if not present
|
|
try:
|
|
import sdnq
|
|
except ImportError:
|
|
print("SDNQ not found. Attempting installation with workaround...")
|
|
try:
|
|
# Try installing from PyPI first (if available)
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", "sdnq", "--no-cache-dir"])
|
|
except:
|
|
print("PyPI install failed. Trying git installation with --no-build-isolation...")
|
|
try:
|
|
# Work around broken pyproject.toml by skipping build isolation
|
|
subprocess.check_call([
|
|
sys.executable, "-m", "pip", "install",
|
|
"git+https://github.com/Disty0/sdnq",
|
|
"--no-build-isolation"
|
|
])
|
|
except Exception as e:
|
|
print(f"SDNQ installation failed: {e}")
|
|
print("MANUAL INSTALLATION REQUIRED:")
|
|
print("1. Clone: git clone https://github.com/Disty0/sdnq")
|
|
print("2. Fix pyproject.toml: change 'license = \"GPL-3.0-only\"' to 'license = {text = \"GPL-3.0-only\"}'")
|
|
print("3. Install: pip install ./sdnq")
|
|
raise ImportError("SDNQ installation failed. Please install manually (see instructions above).")
|
|
|
|
# Try importing again
|
|
try:
|
|
import sdnq
|
|
except ImportError:
|
|
raise ImportError("SDNQ installation failed. Please install manually.")
|
|
|
|
class LoadZImageSDNQ:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_id": ("STRING", {
|
|
"default": "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
|
|
"tooltip": "Hugging Face model ID for the SDNQ-quantized Z-Image Turbo model"
|
|
}),
|
|
"device": (["auto", "cuda", "cpu"], {
|
|
"default": "auto",
|
|
"tooltip": "Device to load the model on. 'auto' selects CUDA if available, otherwise CPU"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"attention_backend": (["default", "flash", "flash3", "sage"], {
|
|
"default": "default",
|
|
"tooltip": "Attention mechanism: 'default' (SDPA), 'flash' (Flash Attention 2), 'flash3' (Flash Attention 3), 'sage' (Sage Attention). Flash/Sage may improve speed on supported GPUs"
|
|
}),
|
|
"enable_compilation": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Compile transformer with torch.compile for faster inference. First run will be slower while compiling"
|
|
}),
|
|
"cpu_offload": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Enable sequential CPU offloading to reduce VRAM usage. Models are moved between CPU and GPU as needed"
|
|
}),
|
|
"low_cpu_mem_usage": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Use memory-efficient loading method. Recommended to keep enabled"
|
|
}),
|
|
"vae_tiling": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Process VAE decode in tiles to reduce VRAM usage. Useful for high resolutions"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("ZIMAGE_SDNQ_PIPELINE",)
|
|
RETURN_NAMES = ("pipeline",)
|
|
FUNCTION = "load_pipeline"
|
|
CATEGORY = "Z-Image (SDNQ)"
|
|
|
|
def load_pipeline(self, model_id, device, attention_backend="default", enable_compilation=False, cpu_offload=False, low_cpu_mem_usage=True, vae_tiling=False):
|
|
print(f"Loading SDNQ Pipeline from {model_id}...")
|
|
|
|
if device == "auto":
|
|
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
else:
|
|
dev = device
|
|
|
|
pipeline = DiffusionPipeline.from_pretrained(
|
|
model_id,
|
|
torch_dtype=torch.bfloat16,
|
|
low_cpu_mem_usage=low_cpu_mem_usage
|
|
)
|
|
|
|
# Apply attention backend
|
|
if attention_backend != "default":
|
|
print(f"Setting attention backend to: {attention_backend}")
|
|
try:
|
|
if attention_backend == "flash":
|
|
pipeline.transformer.set_attention_backend("flash")
|
|
elif attention_backend == "flash3":
|
|
pipeline.transformer.set_attention_backend("_flash_3")
|
|
elif attention_backend == "sage":
|
|
pipeline.transformer.set_attention_backend("sage")
|
|
except Exception as e:
|
|
print(f"Warning: Failed to set attention backend '{attention_backend}': {e}")
|
|
|
|
# Apply CPU offload if requested
|
|
if cpu_offload:
|
|
print("Enabling CPU offload...")
|
|
pipeline.enable_model_cpu_offload()
|
|
else:
|
|
pipeline.to(dev)
|
|
|
|
# Enable VAE tiling if requested (reduces VRAM usage during decode)
|
|
if vae_tiling:
|
|
print("Enabling VAE tiling...")
|
|
try:
|
|
pipeline.enable_vae_tiling()
|
|
except Exception as e:
|
|
print(f"Warning: VAE tiling failed: {e}")
|
|
|
|
# Compile if requested (must be done after moving to device)
|
|
if enable_compilation:
|
|
print("Compiling transformer (first run will be slower)...")
|
|
try:
|
|
pipeline.transformer.compile()
|
|
except Exception as e:
|
|
print(f"Warning: Compilation failed: {e}")
|
|
|
|
print("SDNQ Pipeline loaded successfully.")
|
|
return (pipeline,)
|
|
|
|
# Resolution presets from official Gradio app
|
|
RES_CHOICES = {
|
|
"1024": [
|
|
"1024x1024 ( 1:1 )",
|
|
"1152x896 ( 9:7 )",
|
|
"896x1152 ( 7:9 )",
|
|
"1152x864 ( 4:3 )",
|
|
"864x1152 ( 3:4 )",
|
|
"1248x832 ( 3:2 )",
|
|
"832x1248 ( 2:3 )",
|
|
"1280x720 ( 16:9 )",
|
|
"720x1280 ( 9:16 )",
|
|
"1344x576 ( 21:9 )",
|
|
"576x1344 ( 9:21 )",
|
|
],
|
|
"1280": [
|
|
"1280x1280 ( 1:1 )",
|
|
"1440x1120 ( 9:7 )",
|
|
"1120x1440 ( 7:9 )",
|
|
"1472x1104 ( 4:3 )",
|
|
"1104x1472 ( 3:4 )",
|
|
"1536x1024 ( 3:2 )",
|
|
"1024x1536 ( 2:3 )",
|
|
"1600x896 ( 16:9 )",
|
|
"896x1600 ( 9:16 )",
|
|
"1680x720 ( 21:9 )",
|
|
"720x1680 ( 9:21 )",
|
|
],
|
|
}
|
|
|
|
# Flatten all presets for dropdown
|
|
ALL_RESOLUTIONS = ["custom"]
|
|
for cat_resolutions in RES_CHOICES.values():
|
|
ALL_RESOLUTIONS.extend(cat_resolutions)
|
|
|
|
class ZImageSDNQGenerate:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("ZIMAGE_SDNQ_PIPELINE",),
|
|
"prompt": ("STRING", {
|
|
"default": "a beautiful landscape",
|
|
"multiline": True,
|
|
"tooltip": "Text description of the image to generate. Be descriptive for best results"
|
|
}),
|
|
"resolution_preset": (ALL_RESOLUTIONS, {
|
|
"default": "1024x1024 ( 1:1 )",
|
|
"tooltip": "Select a resolution preset or choose 'custom' to set manual width/height"
|
|
}),
|
|
"num_inference_steps": ("INT", {
|
|
"default": 9,
|
|
"min": 1,
|
|
"max": 100,
|
|
"tooltip": "Number of denoising steps. Z-Image Turbo is optimized for 5-9 steps. More steps = slower but potentially higher quality"
|
|
}),
|
|
"guidance_scale": ("FLOAT", {
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 20.0,
|
|
"step": 0.1,
|
|
"tooltip": "Classifier-free guidance scale. Z-Image Turbo is trained for guidance_scale=0.0 (recommended). Higher values increase prompt adherence but may reduce quality"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 42,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"tooltip": "Random seed for reproducible generation. Same seed + settings = same image"
|
|
}),
|
|
"shift": ("FLOAT", {
|
|
"default": 3.0,
|
|
"min": 1.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"tooltip": "Time shift for FlowMatch scheduler. Controls sampling distribution: lower (1-2) = more noise-end sampling, higher (4-10) = more clean-end sampling. Default 3.0 is balanced"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"custom_width": ("INT", {
|
|
"default": 1024,
|
|
"min": 256,
|
|
"max": 2048,
|
|
"step": 64,
|
|
"tooltip": "Custom output width (only used when resolution_preset is 'custom'). Must be divisible by 64"
|
|
}),
|
|
"custom_height": ("INT", {
|
|
"default": 1024,
|
|
"min": 256,
|
|
"max": 2048,
|
|
"step": 64,
|
|
"tooltip": "Custom output height (only used when resolution_preset is 'custom'). Must be divisible by 64"
|
|
}),
|
|
"enhance_prompt": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Use Z-Image's built-in prompt enhancer for improved semantic understanding and detail. May not be available in all models"
|
|
}),
|
|
"unload_after_generation": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Move models to CPU after generation to free VRAM. Useful for memory-constrained setups"
|
|
}),
|
|
"gc_cuda": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Run Python garbage collection and clear CUDA cache after generation. Helps reclaim memory"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "Z-Image (SDNQ)"
|
|
|
|
def generate(self, pipeline, prompt, resolution_preset, num_inference_steps, guidance_scale, seed, shift, custom_width=1024, custom_height=1024, enhance_prompt=False, unload_after_generation=False, gc_cuda=False):
|
|
import gc
|
|
|
|
# Parse resolution from preset or use custom values
|
|
if resolution_preset == "custom":
|
|
width = custom_width
|
|
height = custom_height
|
|
print(f"Using custom resolution: {width}x{height}")
|
|
else:
|
|
# Parse preset string format: "1024x1024 ( 1:1 )"
|
|
try:
|
|
resolution_str = resolution_preset.split(" ")[0] # Get "1024x1024"
|
|
width, height = map(int, resolution_str.split("x"))
|
|
print(f"Using preset resolution: {width}x{height} from '{resolution_preset}'")
|
|
except Exception as e:
|
|
print(f"Failed to parse resolution preset '{resolution_preset}': {e}. Using 1024x1024")
|
|
width, height = 1024, 1024
|
|
|
|
print(f"Generating image with SDNQ pipeline (shift={shift})...")
|
|
|
|
# Note: Prompt enhancer is not available in SDNQ models (it's a separate API service)
|
|
# Skip the enhancement attempt to avoid console spam
|
|
|
|
generator = torch.manual_seed(seed)
|
|
|
|
# Configure shift on the scheduler before generation
|
|
original_shift = None
|
|
if hasattr(pipeline.scheduler, 'config'):
|
|
try:
|
|
original_shift = pipeline.scheduler.config.shift
|
|
pipeline.scheduler.config.shift = shift
|
|
print(f"Configured scheduler shift: {shift}")
|
|
except Exception as e:
|
|
print(f"Warning: Could not configure shift on scheduler: {e}")
|
|
|
|
# Z-Image uses FlowMatchEulerDiscreteScheduler
|
|
# shift must be configured on the scheduler, not passed to __call__
|
|
image = pipeline(
|
|
prompt=prompt,
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
guidance_scale=guidance_scale,
|
|
generator=generator
|
|
).images[0]
|
|
|
|
# Restore original shift value if we changed it
|
|
if original_shift is not None and hasattr(pipeline.scheduler, 'config'):
|
|
try:
|
|
pipeline.scheduler.config.shift = original_shift
|
|
except:
|
|
pass
|
|
|
|
# Unload models if requested
|
|
if unload_after_generation:
|
|
print("Unloading models to free VRAM...")
|
|
try:
|
|
# Move models to CPU
|
|
if hasattr(pipeline, 'transformer'):
|
|
pipeline.transformer.to('cpu')
|
|
if hasattr(pipeline, 'text_encoder'):
|
|
pipeline.text_encoder.to('cpu')
|
|
if hasattr(pipeline, 'vae'):
|
|
pipeline.vae.to('cpu')
|
|
except Exception as e:
|
|
print(f"Model unloading failed: {e}")
|
|
|
|
# CUDA garbage collection if requested
|
|
if gc_cuda:
|
|
print("Running CUDA garbage collection...")
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
|
|
# Convert PIL to tensor
|
|
image_np = np.array(image)
|
|
image_tensor = torch.from_numpy(image_np).float() / 255.0
|
|
image_tensor = image_tensor.unsqueeze(0)
|
|
|
|
return (image_tensor,)
|