diff --git a/examples/z_image_fun/predict_turbo_t2i_control_2.1.py b/examples/z_image_fun/predict_turbo_t2i_control_2.1.py index dd833eb..ce85113 100644 --- a/examples/z_image_fun/predict_turbo_t2i_control_2.1.py +++ b/examples/z_image_fun/predict_turbo_t2i_control_2.1.py @@ -3,8 +3,10 @@ import sys import numpy as np import torch +import transformers from diffusers import FlowMatchEulerDiscreteScheduler from omegaconf import OmegaConf +from packaging.version import Version from PIL import Image current_file_path = os.path.abspath(__file__) @@ -28,15 +30,6 @@ from videox_fun.utils.utils import (filter_kwargs, get_image, get_image_latent, get_image_to_video_latent, get_video_to_video_latent, save_videos_grid) -import transformers -from packaging.version import Version - -def _dtype_kwargs(dtype): - """`dtype` keyword of `from_pretrained` exists since transformers 4.56 (PR #39782); - older versions use `torch_dtype`.""" - if Version(transformers.__version__) >= Version("4.56"): - return {"dtype": dtype} - return {"torch_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. @@ -146,6 +139,12 @@ if vae_path is not None: tokenizer = AutoTokenizer.from_pretrained( model_name, subfolder="tokenizer" ) + +# `Qwen3ForCausalLM.from_pretrained` renamed `torch_dtype` -> `dtype` in transformers +# 4.56 (PR #39782); pick whichever keyword the installed version accepts. +def _dtype_kwargs(dtype): + return {"dtype": dtype} if Version(transformers.__version__) >= Version("4.56") else {"torch_dtype": dtype} + text_encoder = Qwen3ForCausalLM.from_pretrained( model_name, subfolder="text_encoder", **_dtype_kwargs(weight_dtype), low_cpu_mem_usage=True, diff --git a/scripts/qwenimage21_fun/README_TRAIN.md b/scripts/qwenimage21_fun/README_TRAIN.md index edadaa4..27ab465 100644 --- a/scripts/qwenimage21_fun/README_TRAIN.md +++ b/scripts/qwenimage21_fun/README_TRAIN.md @@ -73,9 +73,6 @@ docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unc > **Qwen-Image 2.1 specific**: the text encoder is a **Qwen3-VL** model, so the environment needs a `transformers` > build that ships the `qwen3_vl` architecture (newer than the base pin in `requirements.txt`). If > `Qwen3VLForConditionalGeneration` / `Qwen3VLProcessor` import as `None`, your `transformers` is too old. -> -> The YOLO object-mask feature (see [2.3](#23-metadatajson-format)) needs `ultralytics`; `yolov8x-seg.pt` -> downloads automatically on first use. --- @@ -136,10 +133,6 @@ The manifest is the standard image metadata JSON plus one extra `control_file_pa > **You only supply the target + control images. You do NOT supply masks.** The inpaint mask is generated on the > fly: > - A random rectangular hole via `get_random_mask` in the collate. -> - Then, on a random ~70% subset of frames, an **irregular object-shaped mask** produced by a **YOLO-seg** -> detector (`ObjectInstanceDetector`), with its edges randomly dilated / eroded / Gaussian-blurred. This mirrors -> `scripts/qwenimage_fun/train_control.py` and gives the model realistic, object-shaped inpaint holes instead of -> only rectangles. > - The masked image fed to the union branch is always `target * (1 - mask)`. > **RGBA note**: the 2.1 VAE reads RGBA. Training images are loaded as RGB and automatically composited over an diff --git a/scripts/qwenimage21_fun/README_TRAIN_zh-CN.md b/scripts/qwenimage21_fun/README_TRAIN_zh-CN.md index 22a2e5f..42a2977 100644 --- a/scripts/qwenimage21_fun/README_TRAIN_zh-CN.md +++ b/scripts/qwenimage21_fun/README_TRAIN_zh-CN.md @@ -71,8 +71,6 @@ docker run -it -p 7860:7860 --network host --gpus all --security-opt seccomp:unc > **Qwen-Image 2.1 特有**:文本编码器是 **Qwen3-VL** 模型,因此环境需要一个包含 `qwen3_vl` 结构的 `transformers` > 版本(比 `requirements.txt` 里的基线更新)。如果 `Qwen3VLForConditionalGeneration` / `Qwen3VLProcessor` 导入为 > `None`,说明你的 `transformers` 太旧。 -> -> YOLO 目标掩膜功能(见 [2.3](#23-metadatajson-格式))需要 `ultralytics`;`yolov8x-seg.pt` 会在首次使用时自动下载。 --- @@ -130,10 +128,7 @@ modelscope download --dataset PAI/X-Fun-Images-Controls-Demo --local_dir ./datas - `type`:图像数据为 `"image"`。 > **你只需提供目标图 + control 图,不需要提供掩膜。** inpaint 掩膜是即时生成的: -> - 先在 collate 中用 `get_random_mask` 生成随机矩形遮挡。 -> - 然后在随机约 70% 的帧上,再由 **YOLO-seg** 检测器(`ObjectInstanceDetector`)生成**不规则的目标形状掩膜**, -> 其边缘随机做膨胀 / 腐蚀 / 高斯模糊。这与 `scripts/qwenimage_fun/train_control.py` 一致,能让模型见到真实、 -> 目标形状的修补空洞,而不只是矩形。 +> - collate 中用 `get_random_mask` 生成随机矩形遮挡。 > - 送进 union 分支的被遮罩图始终是 `target * (1 - mask)`。 > **RGBA 说明**:2.1 VAE 读取 RGBA。训练图以 RGB 载入,在编码前会自动合成到不透明 alpha 通道,因此你不需要提供 RGBA 数据。 diff --git a/scripts/qwenimage21_fun/train_control.py b/scripts/qwenimage21_fun/train_control.py index 9bf02e5..70c1570 100644 --- a/scripts/qwenimage21_fun/train_control.py +++ b/scripts/qwenimage21_fun/train_control.py @@ -14,6 +14,7 @@ # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and + import argparse import gc import logging @@ -30,12 +31,13 @@ import numpy as np import torch import torch.nn.functional as F import torch.utils.checkpoint +import torchvision.transforms.functional as TF import transformers -from accelerate import Accelerator +from accelerate import Accelerator, FullyShardedDataParallelPlugin from accelerate.logging import get_logger from accelerate.state import AcceleratorState from accelerate.utils import ProjectConfiguration, set_seed -from diffusers import FlowMatchEulerDiscreteScheduler +from diffusers import DDIMScheduler, FlowMatchEulerDiscreteScheduler from diffusers.optimization import get_scheduler from diffusers.training_utils import (EMAModel, compute_density_for_timestep_sampling, @@ -46,6 +48,9 @@ from einops import rearrange from omegaconf import OmegaConf from packaging import version from PIL import Image +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullOptimStateDictConfig, FullStateDictConfig, ShardedOptimStateDictConfig, + ShardedStateDictConfig) from torch.utils.data import RandomSampler from torch.utils.tensorboard import SummaryWriter from torchvision import transforms @@ -61,22 +66,23 @@ for project_root in project_roots: from videox_fun.data import (ASPECT_RATIO_512, ASPECT_RATIO_RANDOM_CROP_512, ASPECT_RATIO_RANDOM_CROP_PROB, AspectRatioBatchImageVideoSampler, - ImageVideoControlDataset, ImageVideoSampler, - RandomSampler, get_closest_ratio, get_random_mask) -from videox_fun.models import (AutoencoderKLQwenImage, - Qwen2_5_VLForConditionalGeneration, - Qwen2Tokenizer, - QwenImageControlTransformer2DModel) -from videox_fun.pipeline import QwenImageControlPipeline + ImageVideoControlDataset, + ImageVideoSampler, RandomSampler, + get_closest_ratio, get_random_mask) +from videox_fun.models import (AutoencoderKLQwenImage21, + Qwen3VLForConditionalGeneration, + Qwen3VLProcessor, + QwenImage21ControlTransformer2DModel) +from videox_fun.pipeline import QwenImage21ControlPipeline from videox_fun.utils.discrete_sampler import DiscreteSampling from videox_fun.utils.fsdp_ema import FSDPEMA from videox_fun.utils.tqdm_bar import PauseAwareTqdm from videox_fun.utils.utils import (calculate_dimensions, get_image_latent, - get_image_to_video_latent, save_videos_grid) if is_wandb_available(): - pass + import wandb + def filter_kwargs(cls, kwargs): import inspect @@ -98,16 +104,9 @@ def generate_timestep_with_lognorm(low, high, shape, device="cpu", generator=Non t = 1 / (1 + torch.exp(-u)) * (high - low) + low return torch.clip(t.to(torch.int32), low, high - 1) -def _pack_latents(latents, batch_size, num_channels_latents, height, width, num_frame=None): - if num_frame is None: - latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) - latents = latents.permute(0, 2, 4, 1, 3, 5) - latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) - else: - latents = latents.view(batch_size, num_channels_latents, num_frame, height // 2, 2, width // 2, 2) - latents = latents.permute(0, 2, 3, 5, 1, 4, 6) - latents = latents.reshape(batch_size, num_frame * (height // 2) * (width // 2), num_channels_latents * 4) - return latents +def _pack_latents(latents, batch_size, num_channels_latents, height, width): + # 2.1 consumes latents unpatched (patch_size = 1), so packing is a plain spatial flatten. + return latents.view(batch_size, num_channels_latents, height * width).transpose(1, 2) def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -134,7 +133,58 @@ check_min_version("0.18.0.dev0") logger = get_logger(__name__, log_level="INFO") -def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerator, weight_dtype, global_step): +SYS_PROMPT = "Comprehend and analyze the provided prompt." +# The prompt is built as a raw template string and passed straight to the processor, rather than going +# through apply_chat_template: the two tokenize differently and the checkpoint expects this one. +PROMPT_TEMPLATE_T2I = ( + f"<|im_start|>system\n{SYS_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{{}}<|im_end|>\n" + f"<|im_start|>assistant\n" +) + +def get_prompt_drop_idx(processor): + # Number of leading system-role tokens to drop from the hidden states. Derived from the tokenized + # system message rather than hardcoded, so it tracks the processor's chat template. + sys_message = [{"role": "system", "content": [{"type": "text", "text": SYS_PROMPT}]}] + sys_tokens = processor.apply_chat_template(sys_message, tokenize=True, return_dict=False) + return len(sys_tokens[0]) + +def get_qwen_prompt_embeds(text_encoder, processor, prompt, template, drop_idx, device, weight_dtype): + # Mirrors QwenImage21Pipeline._get_qwen_prompt_embeds for the text-only (t2i) case. + prompt = [" " if not p else p for p in prompt] + prompts = [template.format(t) for t in prompt] + + model_inputs = processor( + text=prompts, padding=True, padding_side="left", return_tensors="pt" + ).to(device) + + # The hidden states have to be read before the vision-language model's final RMSNorm: that is what + # the transformer was trained on. A forward hook returning the module's input neutralizes the norm. + text_model = getattr(text_encoder.model, "language_model", text_encoder.model) + handle = text_model.norm.register_forward_hook(lambda module, args, output: args[0]) + try: + outputs = text_encoder( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + output_hidden_states=True, + ) + finally: + handle.remove() + hidden_states = outputs.hidden_states[-1] + + split_hidden_states = list(_extract_masked_hidden(hidden_states, model_inputs.attention_mask)) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max(e.size(0) for e in split_hidden_states) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + return prompt_embeds.to(dtype=weight_dtype), encoder_attention_mask + +def log_validation(vae, text_encoder, processor, transformer3d, args, accelerator, weight_dtype, global_step): try: is_deepspeed = type(transformer3d).__name__ == 'DeepSpeedEngine' if is_deepspeed: @@ -146,10 +196,10 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato args.pretrained_model_name_or_path, subfolder="scheduler" ) - pipeline = QwenImageControlPipeline( + pipeline = QwenImage21ControlPipeline( vae=vae, text_encoder=text_encoder, - tokenizer=tokenizer, + processor=processor, transformer=accelerator.unwrap_model(transformer3d) if type(transformer3d).__name__ == 'DistributedDataParallel' else transformer3d, scheduler=scheduler, ) @@ -163,26 +213,29 @@ def log_validation(vae, text_encoder, tokenizer, transformer3d, args, accelerato logger.info(f"Rank {accelerator.process_index} using seed: {rank_seed}") for i in range(len(args.validation_prompts)): + # Control validation drives generation from a control image: derive the output resolution from that + # image's aspect ratio (as in scripts/qwenimage_fun/train_control.py) and load it as a single-frame + # (1, 3, h, w) tensor via get_image_latent. control_image = Image.open(args.validation_paths[i]) width, height = control_image.width, control_image.height - width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height) + width, height = calculate_dimensions(args.image_sample_size * args.image_sample_size, width / height) control_image = get_image_latent(control_image, sample_size=(height, width))[:, :, 0] - + sample = pipeline( args.validation_prompts[i], - negative_prompt = "bad detailed", height = height, width = width, generator = generator, - true_cfg_scale = 4.0, num_inference_steps = 20, control_image = control_image, ).images os.makedirs(os.path.join(args.output_dir, "sample"), exist_ok=True) + # 2.1's VAE decodes to RGBA; JPEG cannot store an alpha channel (it raises "cannot write mode RGBA + # as JPEG"), so save the validation preview as PNG -- matching examples/qwenimage21/predict_t2i.py. image = sample[0].save( os.path.join( args.output_dir, - f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.jpg" + f"sample/sample-{global_step}-rank{accelerator.process_index}-image-{i}.png" ) ) @@ -266,7 +319,7 @@ def parse_args(): type=str, default=None, nargs="+", - help=("A set of control videos evaluated every `--validation_epochs` and logged to `--report_to`."), + help=("A set of control images evaluated every `--validation_epochs` and logged to `--report_to`."), ) parser.add_argument( "--output_dir", @@ -491,12 +544,6 @@ def parse_args(): parser.add_argument( "--random_hw_adapt", action="store_true", help="Whether enable random adapt height and width in datasets." ) - parser.add_argument( - "--token_sample_size", - type=int, - default=512, - help="Sample size of the token.", - ) parser.add_argument( "--train_sampling_steps", type=int, @@ -506,7 +553,7 @@ def parse_args(): parser.add_argument( "--image_sample_size", type=int, - default=512, + default=1024, help="Sample size of the image.", ) parser.add_argument( @@ -564,19 +611,11 @@ def parse_args(): parser.add_argument( "--prompt_template_encode", type=str, - default="<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + default=PROMPT_TEMPLATE_T2I, help=( 'The prompt template for text encoder.' ), ) - parser.add_argument( - "--prompt_template_encode_start_idx", - type=int, - default=34, - help=( - 'The start idx for prompt template.' - ), - ) parser.add_argument( "--abnormal_norm_clip_start", type=int, @@ -612,12 +651,6 @@ def parse_args(): default=1.29, help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.", ) - parser.add_argument( - "--guidance_scale", - type=float, - default=3.5, - help="the FLUX.1 dev variant is a guidance distilled model", - ) args = parser.parse_args() env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) @@ -745,9 +778,12 @@ def main(): ) # Get Tokenizer - tokenizer = Qwen2Tokenizer.from_pretrained( - args.pretrained_model_name_or_path, subfolder="tokenizer" + processor = Qwen3VLProcessor.from_pretrained( + args.pretrained_model_name_or_path, subfolder="processor" ) + # 2.1 drops a fixed number of leading system-role tokens from the hidden states; derive it from + # the processor chat template instead of hardcoding, so it tracks template changes. + prompt_drop_idx = get_prompt_drop_idx(processor) def deepspeed_zero_init_disabled_context_manager(): """ @@ -759,8 +795,6 @@ def main(): return [deepspeed_plugin.zero3_init_context_manager(enable=False)] - config = OmegaConf.load(args.config_path) - # Currently Accelerate doesn't know how to handle multiple models under Deepspeed ZeRO stage 3. # For this to work properly all models must be run through `accelerate.prepare`. But accelerate # will try to assign the same optimizer with the same weights to all models during @@ -772,13 +806,12 @@ def main(): # across multiple gpus and only UNet2DConditionModel will get ZeRO sharded. with ContextManagers(deepspeed_zero_init_disabled_context_manager()): # Get Text encoder - text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( args.pretrained_model_name_or_path, subfolder="text_encoder", torch_dtype=weight_dtype ) text_encoder = text_encoder.eval() - # Get Vae - vae = AutoencoderKLQwenImage.from_pretrained( + vae = AutoencoderKLQwenImage21.from_pretrained( args.pretrained_model_name_or_path, subfolder="vae" ).to(weight_dtype) @@ -786,15 +819,18 @@ def main(): latents_mean = (torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1)).to(accelerator.device) latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(accelerator.device) + # Load the training config so the control transformer is built with the extra control kwargs + # (control_layers / control_in_dim) declared in config/qwenimage21/qwenimage21_control.yaml. + config = OmegaConf.load(args.config_path) + # Get Transformer - transformer3d = QwenImageControlTransformer2DModel.from_pretrained( + transformer3d = QwenImage21ControlTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", torch_dtype=weight_dtype, low_cpu_mem_usage=True, transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), ).to(weight_dtype) - # Freeze vae and text_encoder and set transformer3d to trainable vae.requires_grad_(False) text_encoder.requires_grad_(False) @@ -803,7 +839,7 @@ def main(): if args.transformer_path is not None: print(f"From checkpoint: {args.transformer_path}") if args.transformer_path.endswith("safetensors"): - from safetensors.torch import load_file + from safetensors.torch import load_file, safe_open state_dict = load_file(args.transformer_path) else: state_dict = torch.load(args.transformer_path, map_location="cpu") @@ -816,7 +852,7 @@ def main(): if args.vae_path is not None: print(f"From checkpoint: {args.vae_path}") if args.vae_path.endswith("safetensors"): - from safetensors.torch import load_file + from safetensors.torch import load_file, safe_open state_dict = load_file(args.vae_path) else: state_dict = torch.load(args.vae_path, map_location="cpu") @@ -827,8 +863,8 @@ def main(): assert len(u) == 0 # A good trainable modules is showed below now. - # For 3D Patch: trainable_modules = ['ff.net', 'pos_embed', 'attn2', 'proj_out', 'timepositionalencoding', 'h_position', 'w_position'] - # For 2D Patch: trainable_modules = ['ff.net', 'attn2', 'timepositionalencoding', 'h_position', 'w_position'] + # Full finetune: trainable_modules = ['transformer_blocks', 'img_in', 'txt_in', 'time_text_embed', 'modulation', 'norm_out', 'proj_out', 'pos_embed'] + # Partial finetune: trainable_modules = ['transformer_blocks'] transformer3d.train() if accelerator.is_main_process: accelerator.print( @@ -845,11 +881,12 @@ def main(): if zero_stage == 3: raise NotImplementedError("DeepSpeed Zero-3 does not support EMA.") - ema_module = QwenImageControlTransformer2DModel.from_pretrained( + ema_module = QwenImage21ControlTransformer2DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer", torch_dtype=weight_dtype, low_cpu_mem_usage=True, + transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']), ) if args.use_fsdp: # The EMA copy gets the same FSDP wrap as the live model so that @@ -857,7 +894,7 @@ def main(): ema_transformer3d = FSDPEMA(ema_module, source=transformer3d, accelerator=accelerator, fsdp_plugin=fsdp_plugin) else: ema_module = ema_module.to(weight_dtype) - ema_transformer3d = EMAModel(ema_module.parameters(), model_cls=QwenImageControlTransformer2DModel, model_config=ema_module.config) + ema_transformer3d = EMAModel(ema_module.parameters(), model_cls=QwenImage21ControlTransformer2DModel, model_config=ema_module.config) # `accelerate` 0.16.0 will have better support for customized saving if version.parse(accelerate.__version__) >= version.parse("0.16.0"): @@ -904,12 +941,12 @@ def main(): def load_model_hook(models, input_dir): if args.use_ema: ema_path = os.path.join(input_dir, "transformer_ema") - _, ema_kwargs = QwenImageControlTransformer2DModel.load_config(ema_path, return_unused_kwargs=True) - load_model = QwenImageControlTransformer2DModel.from_pretrained( + _, ema_kwargs = QwenImage21ControlTransformer2DModel.load_config(ema_path, return_unused_kwargs=True) + load_model = QwenImage21ControlTransformer2DModel.from_pretrained( input_dir, subfolder="transformer_ema", low_cpu_mem_usage=True, ) - load_model = EMAModel(load_model.parameters(), model_cls=QwenImageControlTransformer2DModel, model_config=load_model.config) + load_model = EMAModel(load_model.parameters(), model_cls=QwenImage21ControlTransformer2DModel, model_config=load_model.config) load_model.load_state_dict(ema_kwargs) ema_transformer3d.load_state_dict(load_model.state_dict()) @@ -921,7 +958,7 @@ def main(): model = models.pop() # load diffusers style into model - load_model = QwenImageControlTransformer2DModel.from_pretrained( + load_model = QwenImage21ControlTransformer2DModel.from_pretrained( input_dir, subfolder="transformer", low_cpu_mem_usage=True, ) @@ -1029,7 +1066,7 @@ def main(): train_dataset = ImageVideoControlDataset( args.train_data_meta, args.train_data_dir, image_sample_size=args.image_sample_size, - enable_bucket=args.enable_bucket, + enable_bucket=args.enable_bucket, enable_inpaint=True, enable_camera_info=False, enable_subject_info=False, @@ -1058,49 +1095,43 @@ def main(): def _create_special_list(length): if length == 1: return [1.0] - first_element = 0.90 - remaining_sum = 1.0 - first_element - other_elements_value = remaining_sum / (length - 1) - return [first_element] + [other_elements_value] * (length - 1) - - MIN_TARGET = 1024 - - if sample_size < MIN_TARGET: - number_list = [1.0] + if length >= 2: + first_element = 0.90 + remaining_sum = 1.0 - first_element + other_elements_value = remaining_sum / (length - 1) + special_list = [first_element] + [other_elements_value] * (length - 1) + return special_list + + if sample_size >= 1536: + number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio + elif sample_size >= 1024: + number_list = [1, 1.25, 1.5, 2] + image_ratio + elif sample_size >= 768: + number_list = [1, 1.25, 1.5] + image_ratio + elif sample_size >= 512: + number_list = [1] + image_ratio else: - max_allowed_ratio = sample_size / MIN_TARGET - base_ratios = [ - 1.0, - 1.1, 1.2, 1.25, 1.33, 1.5, - 1.75, 2.0, 2.25, 2.5, 2.75, - 3.0, 3.5, 4.0, 5.0, 6.0, 8.0 - ] - candidate_ratios = set(base_ratios + list(image_ratio)) - number_list = sorted([r for r in candidate_ratios if 1.0 <= r <= max_allowed_ratio]) - - if not number_list: - number_list = [1.0] + number_list = [1] if all_choices: return number_list - probs = np.array(_create_special_list(len(number_list))) + number_list_prob = np.array(_create_special_list(len(number_list))) if rng is None: - return np.random.choice(number_list, p=probs) + return np.random.choice(number_list, p = number_list_prob) else: - return rng.choice(number_list, p=probs) + return rng.choice(number_list, p = number_list_prob) # Create new output new_examples = {} new_examples["pixel_values"] = [] new_examples["text"] = [] - # Used in Control Mode + # Used in Control mode new_examples["control_pixel_values"] = [] - - # Used in Inpaint mode - new_examples["mask_pixel_values"] = [] - new_examples["mask"] = [] + # Used in Inpaint mode + new_examples["mask_pixel_values"] = [] + new_examples["mask"] = [] # Get downsample ratio in image pixel_value = examples[0]["pixel_values"] @@ -1129,14 +1160,11 @@ def main(): closest_size = [int(x / 32) * 32 for x in closest_size] # 32 = vae_scale_factor(16)*2: keep latent dims even for example in examples: - # To 0~1 - pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() - pixel_values = pixel_values / 255. - - control_pixel_values = torch.from_numpy(example["control_pixel_values"]).permute(0, 3, 1, 2).contiguous() - control_pixel_values = control_pixel_values / 255. - if args.fix_sample_size is not None: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + # Get adapt hw for resize fix_sample_size = list(map(lambda x: int(x), fix_sample_size)) transform = transforms.Compose([ @@ -1145,6 +1173,10 @@ def main(): transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), ]) elif args.random_ratio_crop: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + # Get adapt hw for resize b, c, h, w = pixel_values.size() th, tw = random_sample_size @@ -1161,6 +1193,10 @@ def main(): transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), ]) else: + # To 0~1 + pixel_values = torch.from_numpy(example["pixel_values"]).permute(0, 3, 1, 2).contiguous() + pixel_values = pixel_values / 255. + # Get adapt hw for resize closest_size = list(map(lambda x: int(x), closest_size)) if closest_size[0] / h > closest_size[1] / w: @@ -1173,18 +1209,22 @@ def main(): transforms.CenterCrop(closest_size), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), ]) - - length = int(len(pixel_values) // 2) - new_examples["pixel_values"].append(transform(pixel_values)[length:length + 1]) - new_examples["control_pixel_values"].append(transform(control_pixel_values)[length:length + 1]) - - new_examples["text"].append(example["text"]) - - mask = get_random_mask(new_examples["pixel_values"][-1].size()) - mask_pixel_values = new_examples["pixel_values"][-1] * (1 - mask) + new_examples["pixel_values"].append(transform(pixel_values)) - new_examples["mask_pixel_values"].append(mask_pixel_values[:1]) - new_examples["mask"].append(mask[:1]) + # Control mode: run the SAME spatial transform on the control image, then build the inpaint + # mask / masked-image pair from the transformed target. control_pixel_values is (f, c, h, w); + # get_random_mask returns a (f, 1, h, w) uint8 mask that broadcasts over channels. This mirrors + # scripts/qwenimage_fun/train_control.py's bucket collate. + control_pixel_values = torch.from_numpy(example["control_pixel_values"]).permute(0, 3, 1, 2).contiguous() + control_pixel_values = control_pixel_values / 255. + new_examples["control_pixel_values"].append(transform(control_pixel_values)) + + mask = get_random_mask(new_examples["pixel_values"][-1].size()) + mask_pixel_values = new_examples["pixel_values"][-1] * (1 - mask) + new_examples["mask_pixel_values"].append(mask_pixel_values) + new_examples["mask"].append(mask) + + new_examples["text"].append(example["text"]) # Limit the number of frames to the same new_examples["pixel_values"] = torch.stack([example for example in new_examples["pixel_values"]]) @@ -1194,31 +1234,10 @@ def main(): # Encode prompts when enable_text_encoder_in_dataloader=True if args.enable_text_encoder_in_dataloader: - template = args.prompt_template_encode - drop_idx = args.prompt_template_encode_start_idx - - txt = [template.format(e) for e in batch['text']] - txt_tokens = tokenizer( - txt, max_length=args.tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt" - ).to(accelerator.device) - encoder_hidden_states = text_encoder( - input_ids=txt_tokens.input_ids, - attention_mask=txt_tokens.attention_mask, - output_hidden_states=True, + prompt_embeds, encoder_attention_mask = get_qwen_prompt_embeds( + text_encoder, processor, batch["text"], args.prompt_template_encode, + prompt_drop_idx, accelerator.device, weight_dtype, ) - hidden_states = encoder_hidden_states.hidden_states[-1] - split_hidden_states = _extract_masked_hidden(hidden_states, txt_tokens.attention_mask) - split_hidden_states = [e[drop_idx:] for e in split_hidden_states] - attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] - max_seq_len = max([e.size(0) for e in split_hidden_states]) - prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] - ) - encoder_attention_mask = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] - ) - - prompt_embeds = prompt_embeds.to(dtype=latents.dtype, device=accelerator.device) new_examples['encoder_attention_mask'] = encoder_attention_mask new_examples['encoder_hidden_states'] = prompt_embeds @@ -1268,8 +1287,8 @@ def main(): if fsdp_stage != 0 or zero_stage != 0: from functools import partial - from videox_fun.dist import shard_model - shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.language_model.layers) + from videox_fun.dist import set_multi_gpus_devices, shard_model + shard_fn = partial(shard_model, device_id=accelerator.device, param_dtype=weight_dtype, module_to_wrapper=text_encoder.model.language_model.layers) text_encoder = shard_fn(text_encoder) if args.use_ema: @@ -1359,7 +1378,7 @@ def main(): disable=not accelerator.is_local_main_process, ) - if args.multi_stream and args.train_mode != "normal": + if args.multi_stream: # create extra cuda streams to speedup inpaint vae computation vae_stream_1 = torch.cuda.Stream() vae_stream_2 = torch.cuda.Stream() @@ -1374,34 +1393,18 @@ def main(): batch_sampler.sampler.generator = torch.Generator().manual_seed(args.seed + epoch) for step, batch in enumerate(train_dataloader): # Data batch sanity check - if epoch == first_epoch and step < 1: + if epoch == first_epoch and step == 0: pixel_values, texts = batch['pixel_values'].cpu(), batch['text'] - control_pixel_values = batch["control_pixel_values"].cpu() pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") - control_pixel_values = rearrange(control_pixel_values, "b f c h w -> b c f h w") os.makedirs(os.path.join(args.output_dir, "sanity_check"), exist_ok=True) - for idx, (pixel_value, control_pixel_value, text) in enumerate(zip(pixel_values, control_pixel_values, texts)): + for idx, (pixel_value, text) in enumerate(zip(pixel_values, texts)): pixel_value = pixel_value[None, ...] - control_pixel_value = control_pixel_value[None, ...] gif_name = '-'.join(text.replace('/', '').split()[:10]) if not text == '' else f'{global_step}-{idx}' save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}.gif", rescale=True) - save_videos_grid(control_pixel_value, f"{args.output_dir}/sanity_check/{gif_name[:10]}_control.gif", rescale=True) - - mask_pixel_values, mask, texts = batch['mask_pixel_values'].cpu(), batch['mask'].cpu(), batch['text'] - mask_pixel_values = rearrange(mask_pixel_values, "b f c h w -> b c f h w") - mask = torch.tile(rearrange(mask, "b f c h w -> b c f h w"), [1, 3, 1, 1, 1]) - for idx, (pixel_value, _mask, text) in enumerate(zip(mask_pixel_values, mask, texts)): - pixel_value = pixel_value[None, ...] - _mask = _mask[None, ...] - save_videos_grid(pixel_value, f"{args.output_dir}/sanity_check/mask_pixel_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.gif", rescale=True) - save_videos_grid(_mask, f"{args.output_dir}/sanity_check/mask_{gif_name[:10] if not text == '' else f'{global_step}-{idx}'}.gif", rescale=True) - with accelerator.accumulate(transformer3d): + with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): # Convert images to latent space pixel_values = batch["pixel_values"].to(weight_dtype) - control_pixel_values = batch["control_pixel_values"].to(weight_dtype) - mask_pixel_values = batch["mask_pixel_values"].to(weight_dtype) - mask = batch["mask"].to(weight_dtype) if args.low_vram: torch.cuda.empty_cache() @@ -1410,6 +1413,9 @@ def main(): text_encoder.to("cpu") with torch.no_grad(): + # 2.1's VAE reads RGBA, so composite the RGB batch over an opaque alpha channel. + alpha = torch.ones_like(pixel_values[:, :, :1]) + pixel_values = torch.cat([pixel_values, alpha], dim=2) # This way is quicker when batch grows up def _batch_encode_vae(pixel_values): pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") @@ -1428,24 +1434,42 @@ def main(): else: latents = _batch_encode_vae(pixel_values) latents = ((latents - latents_mean) * latents_std).to(dtype=weight_dtype) - + + # wait for latents = vae.encode(pixel_values) to complete + if vae_stream_1 is not None: + torch.cuda.current_stream().wait_stream(vae_stream_1) + + # --- Control conditioning (control_context) --- + # Encode the control image and the inpaint pair (masked image + latent-resolution mask) and concat them + # into the (b, 129, 1, h', w') tensor the transformer scatters into the joint stream: 64 control-latent + # channels + 1 mask channel + 64 masked-image latent channels. Mirrors + # scripts/qwenimage_fun/train_control.py, adapted to 2.1: the VAE reads RGBA (composite an opaque alpha) + # and t2v_flag scales with 5D-safe broadcasting. Runs after the base latents encode has synced and + # before the VAE is offloaded under --low_vram, so the VAE is never double-booked across CUDA streams. + with torch.no_grad(): + control_pixel_values = batch["control_pixel_values"].to(weight_dtype) + mask_pixel_values = batch["mask_pixel_values"].to(weight_dtype) + mask = batch["mask"].to(weight_dtype) + + control_pixel_values = torch.cat([control_pixel_values, torch.ones_like(control_pixel_values[:, :, :1])], dim=2) control_latents = _batch_encode_vae(control_pixel_values) control_latents = ((control_latents - latents_mean) * latents_std).to(dtype=weight_dtype) + # Drop the control latents 10% of the time so the adapter also sees the unconditional (t2i) path. for bs_index in range(control_latents.size()[0]): - if rng is None: - zero_init_control_conv_in = np.random.choice([0, 1], p = [0.90, 0.10]) - else: - zero_init_control_conv_in = rng.choice([0, 1], p = [0.90, 0.10]) + zero_init_control_conv_in = (np.random.choice([0, 1], p=[0.90, 0.10]) if rng is None + else rng.choice([0, 1], p=[0.90, 0.10])) if zero_init_control_conv_in: control_latents[bs_index] = control_latents[bs_index] * 0 - + + # Downsample the mask to latent resolution. mask is (b, f=1, 1, h, w); squeeze the frame dim so + # interpolate sees a 4D (b, 1, h, w) tensor, then restore a unit frame dim -> (b, 1, 1, h', w'). mask = mask.squeeze(1) - # mask = rearrange(mask, "b f c h w -> b c f h w") - mask_conditions = F.interpolate(1 - mask[:, :1], size=control_latents.size()[-2:], mode='nearest').to(accelerator.device, weight_dtype) + mask_conditions = F.interpolate(1 - mask[:, :1], size=control_latents.size()[-2:], mode='nearest').to(latents.device, weight_dtype) mask_conditions = mask_conditions.unsqueeze(2) - # Encode inpaint latents. + # A full-frame mask carries no inpaint signal, so gate the masked latents off 90% of the time in + # that case (t2v_flag) to stop the model leaning on them. t2v_flag = [(_mask == 1).all() for _mask in mask] new_t2v_flag = [] for _mask in t2v_flag: @@ -1453,19 +1477,16 @@ def main(): new_t2v_flag.append(0) else: new_t2v_flag.append(1) - t2v_flag = torch.from_numpy(np.array(new_t2v_flag)).to(accelerator.device, dtype=weight_dtype) + t2v_flag = torch.from_numpy(np.array(new_t2v_flag)).to(latents.device, dtype=weight_dtype) + mask_pixel_values = torch.cat([mask_pixel_values, torch.ones_like(mask_pixel_values[:, :, :1])], dim=2) mask_latents = _batch_encode_vae(mask_pixel_values) mask_latents = ((mask_latents - latents_mean) * latents_std).to(dtype=weight_dtype) - mask_latents = t2v_flag[:, None, None] * mask_latents + mask_latents = t2v_flag[:, None, None, None, None] * mask_latents - inpaint_latents = torch.concat([mask_conditions, mask_latents], dim=1) + inpaint_latents = torch.cat([mask_conditions, mask_latents], dim=1) control_context = torch.cat([control_latents, inpaint_latents], dim=1) - # wait for latents = vae.encode(pixel_values) to complete - if vae_stream_1 is not None: - torch.cuda.current_stream().wait_stream(vae_stream_1) - if args.low_vram: vae.to('cpu') torch.cuda.empty_cache() @@ -1477,40 +1498,23 @@ def main(): encoder_attention_mask = batch['encoder_attention_mask'] else: with torch.no_grad(): - template = args.prompt_template_encode - drop_idx = args.prompt_template_encode_start_idx - - txt = [template.format(e) for e in batch['text']] - txt_tokens = tokenizer( - txt, max_length=args.tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt" - ).to(accelerator.device) - encoder_hidden_states = text_encoder( - input_ids=txt_tokens.input_ids, - attention_mask=txt_tokens.attention_mask, - output_hidden_states=True, + prompt_embeds, encoder_attention_mask = get_qwen_prompt_embeds( + text_encoder, processor, batch["text"], args.prompt_template_encode, + prompt_drop_idx, accelerator.device, weight_dtype, ) - hidden_states = encoder_hidden_states.hidden_states[-1] - split_hidden_states = _extract_masked_hidden(hidden_states, txt_tokens.attention_mask) - split_hidden_states = [e[drop_idx:] for e in split_hidden_states] - attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] - max_seq_len = max([e.size(0) for e in split_hidden_states]) - prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] - ) - encoder_attention_mask = torch.stack( - [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] - ) - prompt_embeds = prompt_embeds.to(dtype=latents.dtype, device=accelerator.device) - if args.low_vram and not args.enable_text_encoder_in_dataloader: - text_encoder.to('cpu') - torch.cuda.empty_cache() + if args.low_vram and not args.enable_text_encoder_in_dataloader: + text_encoder.to('cpu') + torch.cuda.empty_cache() + with accelerator.accumulate(transformer3d): bsz, channel, num_frame, height, width = latents.size() - latents = _pack_latents(latents, bsz, channel, height, width, num_frame=num_frame) + latents = _pack_latents(latents, bsz, channel, height, width) noise = torch.randn(latents.size(), device=latents.device, generator=torch_rng, dtype=weight_dtype) - control_context = _pack_latents(control_context, bsz, control_context.size(1), height, width, num_frame=num_frame) + # Pack the control conditioning into the same (seq, dim) layout the transformer scatters into the joint + # image positions: (b, 129, 1, h', w') -> (b, h'*w', 129). 2.1 keeps latents unpatched (no num_frame). + control_context = _pack_latents(control_context, bsz, control_context.size(1), height, width) if not args.uniform_sampling: u = compute_density_for_timestep_sampling( @@ -1559,8 +1563,16 @@ def main(): # Add noise target = noise - latents - img_shapes = [[(num_frame, height // 2, width // 2)]] * latents.size(0) - txt_seq_lens = encoder_attention_mask.sum(dim=1).tolist() if encoder_attention_mask is not None else None + # 2.1 keeps latents unpatched, so one img_shapes entry spans the full latent grid, while each + # vision-language image slot in img_mask stands for a 2x2 group of those latent tokens. + img_shapes = [[(1, height, width)]] * latents.size(0) + img_mask = torch.cat( + [ + encoder_attention_mask.new_zeros(bsz, prompt_embeds.size(1), dtype=torch.bool), + encoder_attention_mask.new_ones(bsz, latents.size(1) // 4, dtype=torch.bool), + ], + dim=1, + ) # Predict the noise residual with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device): @@ -1570,11 +1582,11 @@ def main(): encoder_hidden_states_mask=encoder_attention_mask, encoder_hidden_states=prompt_embeds, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, + img_mask=img_mask, control_context=control_context, return_dict=False, - ) - + )[0][:, -noisy_latents.size(1):] + def custom_mse_loss(noise_pred, target, weighting=None, threshold=50): noise_pred = noise_pred.float() target = target.float() @@ -1586,7 +1598,7 @@ def main(): masked_loss = masked_loss * weighting final_loss = masked_loss.mean() return final_loss - + weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas) loss = custom_mse_loss(noise_pred.float(), target.float(), weighting.float()) loss = loss.mean() @@ -1680,7 +1692,7 @@ def main(): log_validation( vae, text_encoder, - tokenizer, + processor, transformer3d, args, accelerator, @@ -1706,7 +1718,7 @@ def main(): log_validation( vae, text_encoder, - tokenizer, + processor, transformer3d, args, accelerator,