import os import sys import numpy as np import torch import onnxruntime as ort from contextlib import contextmanager from omegaconf import OmegaConf from PIL import Image from loguru import logger current_file_path = os.path.abspath(__file__) project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))] for project_root in project_roots: sys.path.insert(0, project_root) if project_root not in sys.path else None repo_root = project_roots[-1] onnx_models_dir = os.path.join(repo_root, "onnx-models") vae_encoder_onnx_path = os.path.join(onnx_models_dir, "vae_encoder_simp_slim.onnx") vae_decoder_onnx_path = os.path.join(onnx_models_dir, "vae_decoder_simp_slim.onnx") # vae_encoder_onnx_path = "" # vae_decoder_onnx_path = "" use_transformer_onnx = True use_controlnet = False # 控制是否启用 ControlNet，默认启用 controlnet_onnx_path = os.path.join(onnx_models_dir, "z_image_controlnet_simp_slim.onnx") transformer_body_onnx_path = os.path.join(onnx_models_dir, "z_image_transformer_body_simp_slim.onnx") onnx_export_dtype = torch.float16 _printed_onnx_debug = False _body_input_cache = {} def _select_onnx_providers(): avail = ort.get_available_providers() if "CUDAExecutionProvider" in avail: return ["CUDAExecutionProvider", "CPUExecutionProvider"] if "AzureExecutionProvider" in avail: return ["AzureExecutionProvider", "CPUExecutionProvider"] return ["CPUExecutionProvider"] onnx_providers = _select_onnx_providers() from diffusers import FlowMatchEulerDiscreteScheduler from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler from videox_fun.models import (AutoencoderKL, AutoTokenizer, Qwen3ForCausalLM, ZImageControlTransformer2DModel) from typing import List, Optional, Union from diffusers.utils.torch_utils import randn_tensor from videox_fun.utils.utils import get_image_latent torch.set_grad_enabled(False) _ort_sessions = {} def _ort_type_to_torch(type_str: str): if type_str is None: return None if "float16" in type_str: return torch.float16 if "bfloat16" in type_str: return torch.bfloat16 if "float" in type_str: return torch.float32 return None def ort_inference(onnx_model_path: str, inputs: dict, providers=None): providers = providers or ["CPUExecutionProvider"] cache_key = (onnx_model_path, tuple(providers)) if cache_key not in _ort_sessions: if not os.path.exists(onnx_model_path): raise FileNotFoundError(f"ONNX model not found: {onnx_model_path}") _ort_sessions[cache_key] = ort.InferenceSession(onnx_model_path, providers=providers) session = _ort_sessions[cache_key] # 按 ONNX 输入类型做 dtype 对齐, bfloat16 等不被 numpy 支持的类型会被安全转换 input_type_map = {i.name: _ort_type_to_torch(getattr(i, "type", None)) for i in session.get_inputs()} inputs_onnx = {} for k, v in inputs.items(): target_torch_dtype = input_type_map.get(k) if isinstance(v, np.ndarray): if target_torch_dtype is not None: np_dtype = np.float16 if target_torch_dtype == torch.float16 else np.float32 if v.dtype != np_dtype: v = v.astype(np_dtype) inputs_onnx[k] = v elif torch.is_tensor(v): if target_torch_dtype is not None and v.dtype != target_torch_dtype: v = v.to(dtype=target_torch_dtype) elif v.dtype == torch.bfloat16: v = v.to(dtype=torch.float32) inputs_onnx[k] = v.detach().to("cpu").numpy() else: raise TypeError(f"Unsupported input type for key {k}: {type(v)}") return session.run(None, inputs_onnx) def _infer_module_device(module: torch.nn.Module) -> torch.device: param = next(module.parameters(), None) if param is not None: return param.device buffer = next(module.buffers(), None) if buffer is not None: return buffer.device return torch.device("cpu") @contextmanager def module_to_device(module: torch.nn.Module, target_device: torch.device): if module is None: yield module return original_device = _infer_module_device(module) target_device = target_device or original_device needs_move = original_device != target_device moved_to_cuda = needs_move and target_device.type == "cuda" if needs_move: module.to(target_device) try: yield module finally: if needs_move: module.to(original_device) if moved_to_cuda and torch.cuda.is_available(): cache_device = target_device.index if cache_device is None: cache_device = torch.cuda.current_device() with torch.cuda.device(cache_device): torch.cuda.empty_cache() # Config and model path config_path = "config/z_image/z_image_control.yaml" model_name = "models/Diffusion_Transformer/Z-Image-Turbo/" # Use torch.float16 if GPU does not support torch.bfloat16 # ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16 weight_dtype = torch.bfloat16 # control_image = "asset/pose.jpg" control_image = "asset/pose_1024x1024.png" control_context_scale = 0.75 device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') # Get tokenizer and text_encoder tokenizer = AutoTokenizer.from_pretrained( model_name, subfolder="tokenizer" ) text_encoder = Qwen3ForCausalLM.from_pretrained( model_name, subfolder="text_encoder", torch_dtype=weight_dtype, low_cpu_mem_usage=True, ) text_encoder.eval() def _encode_prompt( prompt: Union[str, List[str]], device: Optional[torch.device] = None, prompt_embeds: Optional[List[torch.FloatTensor]] = None, max_sequence_length: int = 512, ) -> List[torch.FloatTensor]: device = device or torch.device("cpu") if prompt_embeds is not None: return prompt_embeds if isinstance(prompt, str): prompt = [prompt] for i, prompt_item in enumerate(prompt): messages = [ {"role": "user", "content": prompt_item}, ] prompt_item = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=True, ) prompt[i] = prompt_item text_inputs = tokenizer( prompt, padding="max_length", max_length=max_sequence_length, truncation=True, return_tensors="pt", ) text_input_ids = text_inputs.input_ids.to(device) prompt_masks = text_inputs.attention_mask.to(device).bool() with module_to_device(text_encoder, device): prompt_embeds = text_encoder( input_ids=text_input_ids, attention_mask=prompt_masks, output_hidden_states=True, ).hidden_states[-2] # 固定为 torch.Size([1, 512, 2560] embeddings_list = [] for i in range(len(prompt_embeds)): # embeddings_list.append(prompt_embeds[i][prompt_masks[i]]) embeddings_list.append(prompt_embeds[i]) return embeddings_list def encode_prompt( prompt: Union[str, List[str]], device: Optional[torch.device] = None, do_classifier_free_guidance: bool = True, negative_prompt: Optional[Union[str, List[str]]] = None, prompt_embeds: Optional[List[torch.FloatTensor]] = None, negative_prompt_embeds: Optional[torch.FloatTensor] = None, max_sequence_length: int = 512, ): prompt = [prompt] if isinstance(prompt, str) else prompt prompt_embeds = _encode_prompt( prompt=prompt, device=device, prompt_embeds=prompt_embeds, max_sequence_length=max_sequence_length, ) if do_classifier_free_guidance: if negative_prompt is None: negative_prompt = ["" for _ in prompt] else: negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt assert len(prompt) == len(negative_prompt) negative_prompt_embeds = _encode_prompt( prompt=negative_prompt, device=device, prompt_embeds=negative_prompt_embeds, max_sequence_length=max_sequence_length, ) else: negative_prompt_embeds = [] return prompt_embeds, negative_prompt_embeds def _stack_prompt_embeddings(prompt_embeds_input): if isinstance(prompt_embeds_input, list): return torch.stack(prompt_embeds_input, dim=0) return prompt_embeds_input def _onnx_shape_compatible(model_path: str, providers, latent_shape, prompt_shape, verbose=False): try: sess = ort.InferenceSession(model_path, providers=providers) inputs = {i.name: i for i in sess.get_inputs()} ok = True mismatch_msgs = [] if "latent_model_input" in inputs: shape = inputs["latent_model_input"].shape exp_h, exp_w = shape[3], shape[4] if isinstance(exp_h, int) and exp_h != latent_shape[3]: ok = False mismatch_msgs.append(f"latent_h expected {exp_h}, got {latent_shape[3]}") if isinstance(exp_w, int) and exp_w != latent_shape[4]: ok = False mismatch_msgs.append(f"latent_w expected {exp_w}, got {latent_shape[4]}") if "prompt_embeds" in inputs: pshape = inputs["prompt_embeds"].shape exp_seq = pshape[1] if isinstance(exp_seq, int) and exp_seq != prompt_shape[1]: ok = False mismatch_msgs.append(f"seq_len expected {exp_seq}, got {prompt_shape[1]}") if verbose and (not ok or mismatch_msgs): print(f"[DEBUG] ONNX shape check for {model_path}") print(f" providers={providers}") print(f" model latent shape={inputs.get('latent_model_input').shape if 'latent_model_input' in inputs else 'n/a'}") print(f" model prompt shape={inputs.get('prompt_embeds').shape if 'prompt_embeds' in inputs else 'n/a'}") print(f" runtime latent shape={latent_shape}, prompt shape={prompt_shape}") if mismatch_msgs: print(f" mismatch: {', '.join(mismatch_msgs)}") return ok except Exception as exc: print(f"ONNX shape check failed for {model_path}: {exc}") return True def prepare_latents( batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None, ): height = 2 * (int(height) // (vae_scale_factor * 2)) width = 2 * (int(width) // (vae_scale_factor * 2)) shape = (batch_size, num_channels_latents, height, width) if latents is None: latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) else: if latents.shape != shape: raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") latents = latents.to(device) return latents def _make_zero_control_hints(model_path: str, providers, device, dtype): cache_key = (model_path, tuple(providers)) if cache_key not in _body_input_cache: sess = ort.InferenceSession(model_path, providers=providers) inputs = {i.name: i for i in sess.get_inputs()} if "control_hints" not in inputs: raise RuntimeError("Body ONNX 未找到 control_hints 输入，无法构造零填充。") _body_input_cache[cache_key] = inputs["control_hints"].shape shape = _body_input_cache[cache_key] # 将 None 或字符串符号维度用 1 填充（静态模型应全是 int） resolved = [] for d in shape: if isinstance(d, int): resolved.append(d) else: resolved.append(1) return torch.zeros(resolved, device=device, dtype=dtype) # GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload]. # model_full_load means that the entire model will be moved to the GPU. # # model_full_load_and_qfloat8 means that the entire model will be moved to the GPU, # and the transformer model has been quantized to float8, which can save more GPU memory. # # model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory. # # model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use, # and the transformer model has been quantized to float8, which can save more GPU memory. # # sequential_cpu_offload means that each layer of the model will be moved to the CPU after use, # resulting in slower speeds but saving a large amount of GPU memory. GPU_memory_mode = "model_cpu_offload" # Multi GPUs config # Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used. # For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4. # If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1. ulysses_degree = 1 ring_degree = 1 # Use FSDP to save more GPU memory in multi gpus. fsdp_dit = False fsdp_text_encoder = False # Compile will give a speedup in fixed resolution and need a little GPU memory. # The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload. compile_dit = False # model path # Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++" sampler_name = "Flow" # Load pretrained model if need transformer_path = "models/Personalized_Model/Z-Image-Turbo-Fun-Controlnet-Union.safetensors" vae_path = None lora_path = None # Other params sample_size = [512, 512] # H, W # 使用更长的neg prompt如"模糊, 突变, 变形, 失真, 画面暗, 文本字幕, 画面固定, 连环画, 漫画, 线稿, 没有主体.", 可以增加稳定性 # 在neg prompt中添加"安静, 固定"等词语可以增加动态性. # prompt = "一位年轻女子站在阳光明媚的海岸线上, 白裙在轻拂的海风中微微飘动.她拥有一头鲜艳的紫色长发, 在风中轻盈舞动, 发间系着一个精致的黑色蝴蝶结, 与身后柔和的蔚蓝天空形成鲜明对比.她面容清秀, 眉目精致, 透着一股甜美的青春气息；神情柔和, 略带羞涩, 目光静静地凝望着远方的地平线, 双手自然交叠于身前, 仿佛沉浸在思绪之中.在她身后, 是辽阔无垠、波光粼粼的大海, 阳光洒在海面上, 映出温暖的金色光晕." # prompt = "一位身穿白色仙袍的仙人女子手持青色仙剑, 她抬头望着天空疾驰而来的雷霆, 面容严肃, 一袭紫色长发在风中飘扬, 仿佛与天地间的风雷共舞.她站立在一片古老的山巅, 背后是连绵起伏的群山和翻滚的乌云, 整个场景充满了神秘而壮丽的气息.天空中闪电划过, 照亮了她坚定的眼神和手中的仙剑, 彷佛预示着一场即将到来的大战.她的姿态优雅而坚定, 彷佛是天地间的守护者, 准备迎接任何挑战." prompt = "(masterpiece, best quality, ultra detailed, 8k, CG unity wallpaper),1 young beautiful girl, full body, official art, extremely detailed, highly detailed, 1 girl, aqua eyes, light smile, grey hair, hair flower,bracelet, choker, ribbon, JK, looking at viewer, on the beach, in summer," negative_prompt = " " guidance_scale = 0.00 seed = 43 num_inference_steps = 9 lora_weight = 0.55 save_path = "samples/z-image-t2i-control" config = OmegaConf.load(config_path) transformer = ZImageControlTransformer2DModel.from_pretrained( model_name, subfolder="transformer", low_cpu_mem_usage=True, torch_dtype=weight_dtype, transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), ).to(weight_dtype).to(device) if transformer_path is not None: print(f"From checkpoint: {transformer_path}") if transformer_path.endswith("safetensors"): from safetensors.torch import load_file, safe_open state_dict = load_file(transformer_path) else: state_dict = torch.load(transformer_path, map_location="cpu") state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict m, u = transformer.load_state_dict(state_dict, strict=False) print(f"missing keys: {len(m)}, unexpected keys: {len(u)}") # Get Scheduler Chosen_Scheduler = { "Flow": FlowMatchEulerDiscreteScheduler, "Flow_Unipc": FlowUniPCMultistepScheduler, "Flow_DPM++": FlowDPMSolverMultistepScheduler, }[sampler_name] scheduler = Chosen_Scheduler.from_pretrained( model_name, subfolder="scheduler" ) # pipeline = ZImageControlPipeline( # vae=vae, # tokenizer=tokenizer, # text_encoder=text_encoder, # transformer=transformer, # scheduler=scheduler, # ) # if ulysses_degree > 1 or ring_degree > 1: # from functools import partial # transformer.enable_multi_gpus_inference() # if fsdp_dit: # shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=list(transformer.transformer_blocks) + list(transformer.single_transformer_blocks)) # pipeline.transformer = shard_fn(pipeline.transformer) # print("Add FSDP DIT") # if fsdp_text_encoder: # shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers, ignored_modules=[text_encoder.language_model.embed_tokens], transformer_layer_cls_to_wrap=["MistralDecoderLayer", "PixtralTransformer"]) # text_encoder = shard_fn(text_encoder) # print("Add FSDP TEXT ENCODER") # if compile_dit: # for i in range(len(pipeline.transformer.transformer_blocks)): # pipeline.transformer.transformer_blocks[i] = torch.compile(pipeline.transformer.transformer_blocks[i]) # print("Add Compile") # if GPU_memory_mode == "sequential_cpu_offload": # pipeline.enable_sequential_cpu_offload(device=device) # elif GPU_memory_mode == "model_cpu_offload_and_qfloat8": # convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device) # convert_weight_dtype_wrapper(transformer, weight_dtype) # pipeline.enable_model_cpu_offload(device=device) # elif GPU_memory_mode == "model_cpu_offload": # pipeline.enable_model_cpu_offload(device=device) # elif GPU_memory_mode == "model_full_load_and_qfloat8": # convert_model_weight_to_float8(transformer, exclude_module_name=["img_in", "txt_in", "timestep"], device=device) # convert_weight_dtype_wrapper(transformer, weight_dtype) # pipeline.to(device=device) # else: # pipeline.to(device=device) generator = torch.Generator(device=device).manual_seed(seed) # if lora_path is not None: # pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) # import pdb; pdb.set_trace() if control_image is not None: control_image = get_image_latent(control_image, sample_size=sample_size)[:, :, 0] # torch.Size([1, 3, sample_size[0], sample_size[1]]) print(control_image.shape) height, width = sample_size vae_scale_factor = 8 vae_scale = vae_scale_factor * 2 if height % vae_scale != 0: raise ValueError( f"Height must be divisible by {vae_scale} (got {height}). " f"Please adjust the height to a multiple of {vae_scale}." ) if width % vae_scale != 0: raise ValueError( f"Width must be divisible by {vae_scale} (got {width}). " f"Please adjust the width to a multiple of {vae_scale}." ) _guidance_scale = guidance_scale = 0.0 _joint_attention_kwargs = joint_attention_kwargs = None _interrupt = False _cfg_normalization = cfg_normalization = False _cfg_truncation = cfg_truncation = 1.0 # 2. Define call parameters prompt_embeds = None if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = len(prompt_embeds) batch_size = 1 weight_dtype = text_encoder.dtype # text_encoder.dtype num_channels_latents = 16 # transformer.in_channels inpaint_latent = None # TODO: 原始代码中未定义该变量 from diffusers.image_processor import VaeImageProcessor image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) # Get Vae vae = AutoencoderKL.from_pretrained( model_name, subfolder="vae" ).to(weight_dtype) vae.eval() vae_config_shift_factor = getattr(vae.config, "shift_factor", 0.0) vae_config_scaling_factor = getattr(vae.config, "scaling_factor", 1.0) if control_image is not None: control_image = image_processor.preprocess(control_image, height=height, width=width) control_image = control_image.to(dtype=weight_dtype, device=device) if os.path.exists(vae_encoder_onnx_path): control_latents = ort_inference( vae_encoder_onnx_path, {"pixel_values": control_image}, )[0] control_latents = torch.from_numpy(control_latents).to(device=device, dtype=weight_dtype) else: with module_to_device(vae, device): control_latents = vae.encode(control_image)[0].mode() control_latents = (control_latents - vae_config_shift_factor) * vae_config_scaling_factor else: control_latents = torch.zeros_like(inpaint_latent) control_context = control_latents.unsqueeze(2) control_context_scale_tensor = torch.tensor([control_context_scale], device=device, dtype=torch.float32) if not use_controlnet: control_context = torch.zeros_like(control_context) control_context_scale_tensor = torch.zeros_like(control_context_scale_tensor) onnx_transformer_enabled = use_transformer_onnx and os.path.exists(controlnet_onnx_path) and os.path.exists(transformer_body_onnx_path) if use_transformer_onnx and not onnx_transformer_enabled: print(f"ONNX transformer requested but missing files, fallback to torch. controlnet={controlnet_onnx_path}, body={transformer_body_onnx_path}") else: print(f"[DEBUG] ONNX providers: {onnx_providers}") print(f"[DEBUG] controlnet onnx path: {controlnet_onnx_path}") print(f"[DEBUG] transformer body onnx path: {transformer_body_onnx_path}") do_classifier_free_guidance = False negative_prompt_embeds = None max_sequence_length = 128 # 512 prompt_embeds, negative_prompt_embeds = encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=do_classifier_free_guidance, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, device=device, max_sequence_length=max_sequence_length, ) num_images_per_prompt = 1 latents = None """ (Pdb) latents[0, 0, 100:105, 100] tensor([-0.9203, 1.3958, 0.8130, -0.5280, -1.9788], device='cuda:0') """ # 4. Prepare latent variables latents = prepare_latents( batch_size * num_images_per_prompt, num_channels_latents, height, width, torch.float32, device, generator, latents, ) # torch.Size([1, 16, 216, 124]) print(latents.shape) # 以上代码验证正确 # Repeat prompt_embeds for num_images_per_prompt if num_images_per_prompt > 1: prompt_embeds = [pe for pe in prompt_embeds for _ in range(num_images_per_prompt)] if do_classifier_free_guidance and negative_prompt_embeds: negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] actual_batch_size = batch_size * num_images_per_prompt image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) # Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu # 5. Prepare timesteps mu = calculate_shift( image_seq_len, scheduler.config.get("base_image_seq_len", 256), scheduler.config.get("max_image_seq_len", 4096), scheduler.config.get("base_shift", 0.5), scheduler.config.get("max_shift", 1.15), ) scheduler.sigma_min = 0.0 scheduler_kwargs = {"mu": mu} import inspect # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, sigmas: Optional[List[float]] = None, **kwargs, ): r""" Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. Args: scheduler (`SchedulerMixin`): The scheduler to get timesteps from. num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. timesteps (`List[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. sigmas (`List[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") if timesteps is not None: accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accept_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps sigmas = None timesteps, num_inference_steps = retrieve_timesteps( scheduler, num_inference_steps, device, sigmas=sigmas, **scheduler_kwargs, ) num_warmup_steps = max(len(timesteps) - num_inference_steps * scheduler.order, 0) """ (Pdb) timesteps tensor([1000.0000, 954.5454, 900.0000, 833.3333, 750.0000, 642.8571, 500.0000, 300.0000, 0.0000], device='cuda:0') """ _num_timesteps = len(timesteps) # 9 # 以上代码验证正确 # import pdb; pdb.set_trace() callback_on_step_end = None callback_on_step_end_tensor_inputs = ['latents'] # 6. Denoising loop # with progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]) timestep = (1000 - timestep) / 1000 # Normalized time for time-aware config (0 at start, 1 at end) t_norm = timestep[0].item() # Handle cfg truncation current_guidance_scale = guidance_scale if ( do_classifier_free_guidance and _cfg_truncation is not None and float(_cfg_truncation) <= 1 ): if t_norm > _cfg_truncation: current_guidance_scale = 0.0 # Run CFG only if configured AND scale is non-zero apply_cfg = do_classifier_free_guidance and current_guidance_scale > 0 if apply_cfg: latents_typed = latents.to(transformer.dtype) latent_model_input = latents_typed.repeat(2, 1, 1, 1) prompt_embeds_model_input = [p.to(transformer.dtype) for p in (prompt_embeds + negative_prompt_embeds)] timestep_model_input = timestep.repeat(2) else: latent_model_input = latents.to(transformer.dtype) prompt_embeds_model_input = [p.to(transformer.dtype) for p in prompt_embeds] timestep_model_input = timestep latent_model_input = latent_model_input.unsqueeze(2) prompt_embeds_tensor = _stack_prompt_embeddings(prompt_embeds_model_input) if not _printed_onnx_debug: print(f"[DEBUG] runtime latent shape: {latent_model_input.shape}, dtype={latent_model_input.dtype}") print(f"[DEBUG] runtime prompt shape: {prompt_embeds_tensor.shape}, dtype={prompt_embeds_tensor.dtype}") print(f"[DEBUG] runtime control_context shape: {control_context.shape}, dtype={control_context.dtype}") if onnx_transformer_enabled: onnx_latent = latent_model_input.to(dtype=onnx_export_dtype) onnx_prompt = prompt_embeds_tensor.to(dtype=onnx_export_dtype) onnx_control = control_context.to(dtype=onnx_export_dtype) if apply_cfg: onnx_control = onnx_control.repeat(2, 1, 1, 1, 1) if not _onnx_shape_compatible(controlnet_onnx_path, onnx_providers, onnx_latent.shape, onnx_prompt.shape, verbose=True): print(f"ONNX controlnet 输入尺寸与当前推理不匹配，回退到 torch。模型期望与当前 latent/prompt 尺寸不同。") onnx_transformer_enabled = False if onnx_transformer_enabled and not _onnx_shape_compatible(transformer_body_onnx_path, onnx_providers, onnx_latent.shape, onnx_prompt.shape, verbose=True): print(f"ONNX transformer body 输入尺寸与当前推理不匹配，回退到 torch。模型期望与当前 latent/prompt 尺寸不同。") if onnx_transformer_enabled: control_inputs = { "latent_model_input": onnx_latent, "timestep": timestep_model_input.to(dtype=torch.float32), "prompt_embeds": onnx_prompt, "control_context": onnx_control, } if use_controlnet: control_hints = ort_inference(controlnet_onnx_path, control_inputs, providers=onnx_providers)[0] control_hints_tensor = torch.from_numpy(control_hints).to(device=device, dtype=onnx_export_dtype) else: control_hints_tensor = _make_zero_control_hints(transformer_body_onnx_path, onnx_providers, device, onnx_export_dtype) body_inputs = { "latent_model_input": onnx_latent, "timestep": timestep_model_input.to(dtype=torch.float32), "prompt_embeds": onnx_prompt, "control_hints": control_hints_tensor, "control_context_scale": torch.zeros_like(control_context_scale_tensor).to(dtype=torch.float32) if not use_controlnet else control_context_scale_tensor.to(dtype=torch.float32), } model_out = ort_inference(transformer_body_onnx_path, body_inputs, providers=onnx_providers)[0] model_out_tensor = torch.from_numpy(model_out).to(device=device, dtype=onnx_export_dtype) model_out_list = list(model_out_tensor) else: latent_model_input_list = [t.to(transformer.dtype) for t in latent_model_input.unbind(dim=0)] control_context_input = control_context.repeat(2, 1, 1, 1, 1) if apply_cfg else control_context if not use_controlnet: control_context_input = torch.zeros_like(control_context_input) control_context_scale_zero = torch.zeros_like(control_context_scale_tensor) else: control_context_scale_zero = control_context_scale_tensor control_context_input = control_context_input.to(dtype=transformer.dtype) control_context_scale_zero = control_context_scale_zero.to(dtype=transformer.dtype) # timestep 也保持与权重一致的 dtype，避免 mat1/mat2 dtype 冲突 timestep_model_input = timestep_model_input.to(dtype=transformer.dtype) # import pdb; pdb.set_trace() if device.type == "cuda": with torch.autocast(device_type="cuda", dtype=transformer.dtype): model_out_list = transformer( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, control_context=control_context_input, control_context_scale=control_context_scale_zero, )[0] else: model_out_list = transformer( latent_model_input_list, timestep_model_input, prompt_embeds_model_input, control_context=control_context_input, control_context_scale=control_context_scale_zero, )[0] if not _printed_onnx_debug: _printed_onnx_debug = True if apply_cfg: # Perform CFG pos_out = model_out_list[:actual_batch_size] neg_out = model_out_list[actual_batch_size:] noise_pred = [] for j in range(actual_batch_size): pos = pos_out[j].float() neg = neg_out[j].float() pred = pos + current_guidance_scale * (pos - neg) # Renormalization if _cfg_normalization and float(_cfg_normalization) > 0.0: ori_pos_norm = torch.linalg.vector_norm(pos) new_pos_norm = torch.linalg.vector_norm(pred) max_new_norm = ori_pos_norm * float(_cfg_normalization) if new_pos_norm > max_new_norm: pred = pred * (max_new_norm / new_pos_norm) noise_pred.append(pred) noise_pred = torch.stack(noise_pred, dim=0) else: noise_pred = torch.stack([t.float() for t in model_out_list], dim=0) noise_pred = noise_pred.squeeze(2) noise_pred = -noise_pred # compute the previous noisy sample x_t -> x_t-1 latents = scheduler.step(noise_pred.to(torch.float32), t, latents, return_dict=False)[0] assert latents.dtype == torch.float32 if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) # call the callback, if provided # if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0): # progress_bar.update() # 显存优化, 80G -> 20G """ # 原始数据 (vae onnx 版本) (Pdb) latents.shape torch.Size([1, 16, 216, 124]) (Pdb) latents[0, 0, 100, 100:105] tensor([ 0.8545, 1.0117, 0.7908, -0.7002, 0.3965], device='cuda:0') (Pdb) # 当前数据 launcher.py (Pdb) latents.shape torch.Size([1, 16, 216, 124]) (Pdb) latents[0, 0, 100, 100:105] tensor([ 0.2899, 0.7049, 0.4407, -1.2531, -0.0161], device='cuda:2') (Pdb) # 对齐 embedding 后的结果(torch版本) (Pdb) latents.shape torch.Size([1, 16, 216, 124]) (Pdb) latents[0, 0, 100, 100:105] tensor([ 0.8763, 1.0138, 0.8031, -0.6851, 0.3751], device='cuda:2') # cuda:0 结果相同 (Pdb) # 修改为 onnx 版本 (Pdb) latents[0, 0, 100, 100:105] tensor([ 0.8545, 1.0117, 0.7908, -0.7002, 0.3965], device='cuda:0') # 以上结果对齐 """ output_type = "pil" if output_type == "latent": image = latents else: latents = latents.to(vae.dtype) latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor # torch.Size([1, 16, 216, 124]) if os.path.exists(vae_decoder_onnx_path): image = ort_inference( vae_decoder_onnx_path, {"latents": latents}, )[0] image = torch.from_numpy(image).to(device=device, dtype=vae.dtype) else: with module_to_device(vae, device): image = vae.decode(latents, return_dict=False)[0] # torch.Size([1, 3, 1728, 992]) """ (Pdb) latents[0, 0, 100, 100:105] tensor([ 2.4844, 2.9062, 2.2969, -1.8203, 1.2188], device='cuda:0', dtype=torch.bfloat16) (Pdb) image[0, 0, 100, 100:105] tensor([0.3906, 0.3848, 0.3809, 0.3809, 0.3848], device='cuda:0', dtype=torch.bfloat16) """ # image = ort_inference( # "/data/tmp/yongqiang/nfs/Z-Image-Turbo.axera/python/VideoX-Fun/onnx-models/vae_decoder_slim.onnx", # {"latent": latents}, # )[0] # image = torch.from_numpy(image).to(device=device, dtype=vae.dtype) image = image_processor.postprocess(image, output_type=output_type) # Offload all models # maybe_free_model_hooks() # return_dict = True # if not return_dict: # return (image,) # return ZImagePipelineOutput(images=image) # with torch.no_grad(): # if control_image is not None: # control_image = get_image_latent(control_image, sample_size=sample_size)[:, :, 0] # torch.Size([1, 3, sample_size[0], sample_size[1]]) # sample = pipeline( # prompt = prompt, # negative_prompt = negative_prompt, # height = sample_size[0], # width = sample_size[1], # generator = generator, # guidance_scale = guidance_scale, # control_image = control_image, # num_inference_steps = num_inference_steps, # control_context_scale = control_context_scale, # ).images # if lora_path is not None: # pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype) sample = image def save_results(): if not os.path.exists(save_path): os.makedirs(save_path, exist_ok=True) index = len([path for path in os.listdir(save_path)]) + 1 prefix = str(index).zfill(8) video_path = os.path.join(save_path, prefix + ".png") image = sample[0] image.save(video_path) # if ulysses_degree * ring_degree > 1: # import torch.distributed as dist # if dist.get_rank() == 0: # save_results() # else: # save_results() save_results() logger.info(f"Saved image to {save_path}")