84 lines
3.5 KiB
Python
84 lines
3.5 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from transformers import AutoTokenizer
|
|
from diffusers.schedulers import DDIMScheduler
|
|
|
|
class ZImageTurboQDiTGenerateUnload:
|
|
"""
|
|
Offline generator that takes (MODEL, VAE, TEXT_ENCODER) and runs a DDIM sampling loop.
|
|
Assumes transformer forward signature similar to UNet2DConditionModel; if you're using
|
|
QwenImageTransformer2DModel, this may still work depending on the forward signature.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"vae": ("VAE",),
|
|
"text_encoder": ("TEXT_ENCODER",),
|
|
"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 (Offline)"
|
|
|
|
def generate(self, model, vae, text_encoder, prompt, negative_prompt, height, width, steps, guidance_scale, seed):
|
|
device = next(model.parameters()).device
|
|
torch.manual_seed(seed)
|
|
|
|
# Tokenizer (Qwen tokenizer; if not available, use a generic AutoTokenizer from Qwen repo)
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B") # local cache if available
|
|
except Exception:
|
|
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-7B", local_files_only=False)
|
|
|
|
# Encode prompts
|
|
pos_inputs = tokenizer([prompt], return_tensors="pt", padding=True).to(device)
|
|
neg_inputs = tokenizer([negative_prompt], return_tensors="pt", padding=True).to(device)
|
|
|
|
pos_embeds = text_encoder(**pos_inputs).last_hidden_state
|
|
neg_embeds = text_encoder(**neg_inputs).last_hidden_state
|
|
|
|
# Scheduler
|
|
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
|
|
scheduler.set_timesteps(steps, device=device)
|
|
|
|
# Latent size from image size (standard 8x compression)
|
|
latent_h = height // 8
|
|
latent_w = width // 8
|
|
latents = torch.randn((1, getattr(model, "in_channels", 4), latent_h, latent_w), device=device)
|
|
|
|
# DDIM loop with classifier-free guidance
|
|
for t in scheduler.timesteps:
|
|
latent_input = latents
|
|
# unconditional
|
|
noise_uncond = model(latent_input, t, encoder_hidden_states=neg_embeds).sample
|
|
# conditional
|
|
noise_text = model(latent_input, t, encoder_hidden_states=pos_embeds).sample
|
|
# CFG
|
|
noise_pred = noise_uncond + guidance_scale * (noise_text - noise_uncond)
|
|
|
|
latents = scheduler.step(noise_pred, t, latents).prev_sample
|
|
|
|
# Decode with VAE
|
|
# SD convention: scale by 1/0.18215 if VAE expects normalized latents.
|
|
try:
|
|
image = vae.decode(latents / 0.18215).sample
|
|
except Exception:
|
|
image = vae.decode(latents).sample
|
|
|
|
image = (image.clamp(-1, 1) + 1) / 2 # [-1,1] -> [0,1]
|
|
image = (image * 255).byte().cpu().numpy()
|
|
pil_img = Image.fromarray(image[0].transpose(1, 2, 0))
|
|
return (np.array(pil_img),)
|
|
|