diff --git a/README.md b/README.md index 4809e40..14823bc 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Qwen/Qwen-Image](https://huggingface.co/Qwen/Qwen-Image) (Qwen-Image) - [Qwen/Qwen-Image-2512](https://huggingface.co/Qwen/Qwen-Image-2512) (Qwen-Image-2512) - [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) (Qwen-Image-2.1) +- [Qwen/Qwen-Image-2.1-Turbo](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo) (Qwen-Image-2.1 Turbo) - [HiDream-ai/HiDream-I1-Full](https://huggingface.co/HiDream-ai/HiDream-I1-Full) (HiDream I1) - [OmniGen2/OmniGen2](https://huggingface.co/OmniGen2/OmniGen2) (OmniGen2) - [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) (Z-Image Turbo) @@ -45,6 +46,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [Qwen/Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) (Qwen-Image-Edit-2509) - [Qwen/Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) (Qwen-Image-Edit-2511) - [Qwen/Qwen-Image-2.1](https://huggingface.co/Qwen/Qwen-Image-2.1) (Qwen-Image-2.1) - one model for both; it edits when your dataset has control images +- [Qwen/Qwen-Image-2.1-Turbo](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo) (Qwen-Image-2.1 Turbo) - same, with the 8-step turbo weights - [HiDream-ai/HiDream-E1-1](https://huggingface.co/HiDream-ai/HiDream-E1-1) (HiDream E1) - [Boogu/Boogu-Image-0.1-Edit](https://huggingface.co/Boogu/Boogu-Image-0.1-Edit) (Boogu Image Edit) - [krea/Krea-2-Raw](https://huggingface.co/krea/Krea-2-Raw) (Krea 2 Edit Training) diff --git a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py b/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py index 32e0770..d9cb84b 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py +++ b/extensions_built_in/diffusion_models/qwen_image_2/qwen_image_2.py @@ -19,19 +19,30 @@ see the sampled timestep. Flow-matching convention matches ai-toolkit (t=1 noise -> t=0 clean, target = noise - clean), so `get_noise_prediction` does no time flip or negation. + +Qwen-Image-2.1-Turbo is the same architecture distilled to 8 steps at CFG 1. +Point `name_or_path` at its weight file in the repack +(`Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors`); +`model_kwargs.turbo` samples with the schedule the checkpoint ships. Training +it goes through a training adapter +(`assistant_lora_path`): a LoRA that is live on the DiT while training and +off while sampling, so samples show the turbo model plus what was trained. """ import os from typing import TYPE_CHECKING, List, Optional +import huggingface_hub import numpy as np import torch from PIL import Image +from safetensors.torch import load_file from toolkit.accelerator import unwrap_model from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds from toolkit.basic import flush -from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.config_modules import GenerateImageConfig, ModelConfig, NetworkConfig +from toolkit.lora_special import LoRASpecialNetwork from toolkit.metadata import get_meta_for_safetensors from toolkit.models.base_model import BaseModel from toolkit.samplers.custom_flowmatch_sampler import ( @@ -80,6 +91,16 @@ scheduler_config = { COMFY_REPO = "Comfy-Org/Qwen-Image-2.1" BASE_REPO = "Qwen/Qwen-Image-2.1" +# the turbo checkpoint is sampled with the schedule it ships +# (TURBO_SAMPLE_SIGMAS in src/pipeline.py) as-is: no resolution shift, no +# terminal stretch; the step count is fixed +turbo_scheduler_config = { + **scheduler_config, + "shift": 1.0, + "shift_terminal": None, + "use_dynamic_shifting": False, +} + # decode above this many output pixels goes through the VAE's tiled path TILE_DECODE_ABOVE_PIXELS = 1024 * 1024 @@ -141,6 +162,15 @@ class QwenImage2Model(BaseModel): def get_train_scheduler(): return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) + @property + def is_turbo(self) -> bool: + return bool(self.model_config.model_kwargs.get("turbo", False)) + + def get_sample_scheduler(self): + if self.is_turbo: + return CustomFlowMatchEulerDiscreteScheduler(**turbo_scheduler_config) + return self.get_train_scheduler() + def get_bucket_divisibility(self): # 16 for the VAE, 2 more because the DiT groups target latent tokens # into 2x2 blocks, one per vision slot @@ -165,12 +195,35 @@ class QwenImage2Model(BaseModel): # a local full checkpoint supplies its own text encoder / vae base_model_path = model_path + load_kwargs = self.component_load_kwargs("transformer") self.print_and_status_update("Loading transformer") - transformer = QwenImage21Transformer2DModel.load( - model_path, - config_path=base_model_path, - **self.component_load_kwargs("transformer"), - ) + if self.model_config.assistant_lora_path is None: + transformer = QwenImage21Transformer2DModel.load( + model_path, config_path=base_model_path, **load_kwargs + ) + else: + # the adapter stays a live LoRA on the quantized linears, attached + # after quantization (it must wrap the quantized forward) and + # before layer offloading (the offloader routes through it) + offload = load_kwargs.pop("offload", 0.0) + transformer = QwenImage21Transformer2DModel.load_model( + model_path, + dtype=load_kwargs["dtype"], + config_path=base_model_path, + qtype=(load_kwargs.get("qtype") or "").split("|", 1)[0] or None, + quantize_on_load=False, + ) + transformer.aitk_post_load(offload=0.0, **load_kwargs) + self.load_training_adapter(transformer) + if offload and offload > 0: + from toolkit.memory_management import MemoryManager + + MemoryManager.attach( + transformer, + torch.device(load_kwargs["quantize_device"]), + offload_percent=offload, + ignore_modules=list(transformer.get_offload_ignore_modules() or []), + ) flush() self.print_and_status_update("Loading text encoder") @@ -204,6 +257,66 @@ class QwenImage2Model(BaseModel): self.pipeline = QwenImage21Pipeline(self) self.print_and_status_update("Model Loaded") + def load_training_adapter(self, transformer: QwenImage21Transformer2DModel): + """Attach the training adapter (a LoRA on the DiT) as a live assistant: + on at 1.0 while training, off while sampling so samples show the turbo + model. It is never merged: the weights are int8, whose grid swallows a + delta this small on a requantize.""" + self.print_and_status_update("Loading assistant LoRA") + lora_path = self.model_config.assistant_lora_path + if not os.path.exists(lora_path): + # assume it is a hub path + lora_splits = lora_path.split("/") + if len(lora_splits) != 3: + raise ValueError( + f"Assistant LoRA path {lora_path} is not a valid local path or hub path." + ) + repo_id = "/".join(lora_splits[:2]) + filename = lora_splits[2] + try: + lora_path = huggingface_hub.hf_hub_download(repo_id=repo_id, filename=filename) + self.model_config.assistant_lora_path = lora_path + except Exception as e: + raise ValueError(f"Failed to download assistant LoRA from {lora_path}: {e}") + lora_state_dict = load_file(lora_path) + dim_key = next(k for k in lora_state_dict if k.endswith("lora_A.weight")) + dim = int(lora_state_dict[dim_key].shape[0]) + lora_state_dict = { + key.replace("diffusion_model.", "transformer."): value + for key, value in lora_state_dict.items() + } + + network_config = NetworkConfig(type="lora", linear=dim, linear_alpha=dim, transformer_only=True) + LoRASpecialNetwork.LORA_PREFIX_UNET = "lora_transformer" + network = LoRASpecialNetwork( + text_encoder=None, + unet=transformer, + lora_dim=network_config.linear, + multiplier=1.0, + alpha=network_config.linear_alpha, + train_unet=True, + train_text_encoder=False, + network_config=network_config, + network_type=network_config.type, + transformer_only=network_config.transformer_only, + is_transformer=True, + target_lin_modules=self.target_lora_modules, + is_assistant_adapter=True, + is_ara=True, + # transformer_only filters to the block list only through the holder + base_model=self, + ) + network.apply_to(None, transformer, apply_text_encoder=False, apply_unet=True) + network.force_to(self.device_torch, dtype=self.torch_dtype) + network.eval() + network.requires_grad_(False) + network._update_torch_multiplier() + network.load_weights(lora_state_dict) + self.assistant_lora: LoRASpecialNetwork = network + self.assistant_lora.multiplier = 1.0 + self.assistant_lora.is_active = True + self.invert_assistant_lora = False + # ------------------------------------------------------------------ # VAE. The latents are RGBA. Images without alpha get an opaque one on # encode, and decode drops it again unless RGBA output is on. diff --git a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py b/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py index 710a554..0bd1d51 100644 --- a/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py +++ b/extensions_built_in/diffusion_models/qwen_image_2/src/pipeline.py @@ -303,6 +303,21 @@ def run_transformer( return unpack_latents(output[:, -target_tokens:], height, width) +# Qwen-Image-2.1-Turbo's 8-step schedule, from `sample_sigmas` in the +# checkpoint's model_index.json (Qwen/Qwen-Image-2.1-Turbo); the final step +# runs from the last sigma to 0 +TURBO_SAMPLE_SIGMAS = [ + 1.0, + 0.978453, + 0.95418, + 0.926626, + 0.89508, + 0.845148, + 0.704534, + 0.414568, +] + + def calculate_shift( image_seq_len: int, base_seq_len: int = 256, @@ -367,16 +382,22 @@ class QwenImage21Pipeline: ) latents = latents.to(device, dtype=dtype) - scheduler = model.get_train_scheduler() - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) - mu = calculate_shift( - latent_height * latent_width, - scheduler.config.get("base_image_seq_len", 256), - scheduler.config.get("max_image_seq_len", 8192), - scheduler.config.get("base_shift", 0.5), - scheduler.config.get("max_shift", 0.9), - ) - scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu) + scheduler = model.get_sample_scheduler() + if model.is_turbo: + # the turbo checkpoint's own 8-step schedule; num_inference_steps + # does not override it (other schedules are unevaluated) + sigmas = np.array(TURBO_SAMPLE_SIGMAS, dtype=np.float32) + scheduler.set_timesteps(sigmas=sigmas, device=device) + else: + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + mu = calculate_shift( + latent_height * latent_width, + scheduler.config.get("base_image_seq_len", 256), + scheduler.config.get("max_image_seq_len", 8192), + scheduler.config.get("base_shift", 0.5), + scheduler.config.get("max_shift", 0.9), + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu) scheduler.set_begin_index(0) # 2.1 is meant to be sampled without guidance; a scale of 1 skips it diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index fb33d5a..f270a5d 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -2,7 +2,7 @@ // registers in AI_TOOLKIT_MODELS. Loaded at runtime by the UI, not bundled: // see ui/src/extensions/README.md for the convention and the allowed imports. import Link from "next/link"; -import type { ModelArch } from "@/app/jobs/new/options"; +import type { CustomModelOption, ModelArch } from "@/app/jobs/new/options"; import type { JobConfig } from "@/types"; import { defaultSampleConfig, @@ -12,6 +12,55 @@ import { const defaultNameOrPath = ""; const defaultLinearRank = 32; +// Qwen-Image-2.1 and its turbo distill share the RGBA VAE, so both entries +// carry this option +const qwenImage21TransparencyOption: CustomModelOption[] = [ + { + type: "checkbox", + label: "Transparency (RGBA)", + getValue: (config: JobConfig) => + config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false, + onChange: ( + value: boolean, + config: JobConfig, + setJobConfig: (value: any, key: string) => void, + ) => { + const kwargs = { + ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}), + }; + if (value) { + kwargs.rgba = true; + } else { + delete kwargs.rgba; + } + setJobConfig(kwargs, "config.process[0].model.model_kwargs"); + }, + doc: { + title: "Transparency (RGBA)", + description: ( +
+

+ This model's VAE is natively RGBA, so alpha can be carried + end to end. When on, dataset images and reference images load + with their alpha channel, the VAE encodes all four channels, and + samples are saved as PNGs with their transparency intact. +

+

+ Images with no alpha of their own get a fully opaque one, so a + mixed dataset is fine. Turn this off to train and sample flat + RGB: alpha is dropped on the way in and added back as opaque. +

+

+ Changing this re-caches latents, and it cannot be combined with + a dataset using alpha_mask, which consumes the + alpha channel as a loss mask instead. +

+
+ ), + }, + }, +]; + export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ { name: "anima", @@ -573,52 +622,56 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "model.low_vram", "model.layer_offloading", ], - customModelSelectOptions: [ - { - type: "checkbox", - label: "Transparency (RGBA)", - getValue: (config: JobConfig) => - config?.config?.process?.[0]?.model?.model_kwargs?.rgba ?? false, - onChange: ( - value: boolean, - config: JobConfig, - setJobConfig: (value: any, key: string) => void, - ) => { - const kwargs = { - ...(config?.config?.process?.[0]?.model?.model_kwargs ?? {}), - }; - if (value) { - kwargs.rgba = true; - } else { - delete kwargs.rgba; - } - setJobConfig(kwargs, "config.process[0].model.model_kwargs"); + customModelSelectOptions: qwenImage21TransparencyOption, + }, + { + name: "qwen_image_2:turbo", + label: "Qwen-Image-2.1 Turbo (w/ Training Adapter)", + generateNameOverride: "Qwen-Image-2.1 Turbo", + group: "image", + defaults: { + // default updates when [selected, unselected] in the UI + // the turbo weight file in the repack (text encoder and VAE come from + // the same repack); model_kwargs.turbo selects the turbo schedule + "config.process[0].model.name_or_path": [ + "Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors", + defaultNameOrPath, + ], + "config.process[0].model.quantize": [true, false], + "config.process[0].model.quantize_te": [true, false], + "config.process[0].model.low_vram": [true, false], + "config.process[0].train.unload_text_encoder": [false, false], + "config.process[0].sample.sampler": ["flowmatch", "flowmatch"], + "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], + "config.process[0].train.timestep_type": ["shift", "sigmoid"], + "config.process[0].model.qtype": ["convrot8", "qfloat8"], + "config.process[0].model.qtype_te": ["convrot8", "qfloat8"], + // training adapter: live on the DiT while training, off while sampling + "config.process[0].model.assistant_lora_path": [ + 'ostris/qwen_image_2_turbo_training_adapter/qwen_image_2_turbo_training_adapter_v1.safetensors', + undefined + ], + // the turbo checkpoint samples with its own fixed 8-step schedule at + // CFG 1; the step count here is informational + "config.process[0].sample.guidance_scale": [1.0, 4.0], + "config.process[0].sample.sample_steps": [8, 25], + "config.process[0].model.model_kwargs": [ + { + turbo: true, + rgba: false, }, - doc: { - title: "Transparency (RGBA)", - description: ( -
-

- This model's VAE is natively RGBA, so alpha can be carried - end to end. When on, dataset images and reference images load - with their alpha channel, the VAE encodes all four channels, and - samples are saved as PNGs with their transparency intact. -

-

- Images with no alpha of their own get a fully opaque one, so a - mixed dataset is fine. Turn this off to train and sample flat - RGB: alpha is dropped on the way in and added back as opaque. -

-

- Changing this re-caches latents, and it cannot be combined with - a dataset using alpha_mask, which consumes the - alpha channel as a loss mask instead. -

-
- ), - }, - }, + {}, + ], + }, + disableSections: ["network.conv", "train.unload_text_encoder"], + additionalSections: [ + "datasets.multi_control_paths", + "sample.multi_ctrl_imgs", + "model.low_vram", + "model.layer_offloading", + "model.assistant_lora_path", ], + customModelSelectOptions: qwenImage21TransparencyOption, }, { name: "ming_image", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 011ccc7..4d6e585 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -189,6 +189,18 @@ MODEL_TESTS = { "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 1.0}, "needs_control_image": True, }, + "qwen_image_2:turbo": { + # model_kwargs.turbo samples with the checkpoint's fixed 8-step + # schedule at CFG 1 + "model": { + "name_or_path": "Comfy-Org/Qwen-Image-2.1/diffusion_models/qwen_image_2.1_turbo_int8_convrot.safetensors", + "quantize": True, + "quantize_te": True, + "model_kwargs": {"turbo": True}, + }, + "sample": {**IMG, "num_inference_steps": 8, "guidance_scale": 1.0}, + "needs_control_image": True, + }, "ming_image": { "model": { "name_or_path": "Comfy-Org/Ming-Image",