github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add support for fine-tuning Hidream O1 (#831)

Browse Source 
* Initial support for hidream. Lora keys likely need work

* Fix saving for hidream-o1

* Remove dependence on flash attention for hidream o1

* Fix gradient checkpointing for hidream o1

* A lot of fixes for hidream. Handle loading and saving as comfy model.

* Omit layers not used in comfy. Fix issue with lora loading keys in comfy

* Version bumpo
This commit is contained in:
Jaret Burkett (Ostris)
2026-05-12 11:15:16 -06:00
committed by GitHub
parent f38de2a2fe
commit fc4d6ebf39
14 changed files with 3633 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
extensions_built_in/diffusion_models/__init__.py
View File

@@ -11,6 +11,7 @@ from .ltx2 import LTX2Model, LTX23Model
from .zeta_chroma import ZetaChromaModel
from .ernie_image import ErnieImageModel
from .nucleus_image import NucleusImageModel
from .hidream.hidream_o1_model import HidreamO1Model
AI_TOOLKIT_MODELS = [
# put a list of models here
@@ -36,4 +37,5 @@ AI_TOOLKIT_MODELS = [
ZetaChromaModel,
ErnieImageModel,
NucleusImageModel,
HidreamO1Model,
]

3
extensions_built_in/diffusion_models/hidream/__init__.py
View File

@@ -1,2 +1,3 @@
from .hidream_model import HidreamModel
from .hidream_e1_model import HidreamE1Model
from .hidream_e1_model import HidreamE1Model
from .hidream_o1_model import HidreamO1Model

544
extensions_built_in/diffusion_models/hidream/hidream_o1_model.py Normal file
View File

@@ -0,0 +1,544 @@
import os
from typing import List, Optional
import torch
import yaml
from toolkit.config_modules import GenerateImageConfig, ModelConfig
from toolkit.metadata import get_meta_for_safetensors
from toolkit.models.base_model import BaseModel
from toolkit.basic import flush
from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
from toolkit.prompt_utils import PromptEmbeds
from toolkit.samplers.custom_flowmatch_sampler import (
CustomFlowMatchEulerDiscreteScheduler,
)
from safetensors.torch import load_file, save_file
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from toolkit.memory_management import MemoryManager
from transformers import AutoProcessor
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
from .src.hidream_o1.qwen3_vl_transformers import Qwen3VLForConditionalGeneration
from .src.hidream_o1.pipeline import HiDreamO1Pipeline, DEFAULT_NOISE_SCALE
from toolkit.models.FakeVAE import FakeVAE
from typing import TYPE_CHECKING
from .src.hidream_o1.model_config import model_config
if TYPE_CHECKING:
from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
scheduler_config = {
"num_train_timesteps": 1000,
"shift": 3.0,
"use_dynamic_shifting": False,
}
_GLOBAL_NOISE_SCALE = DEFAULT_NOISE_SCALE
class HidreamO1FlowmatchScheduler(CustomFlowMatchEulerDiscreteScheduler):
def __init__(self, *args, **kwargs):
self.noise_scale = kwargs.get("noise_scale", DEFAULT_NOISE_SCALE)
# remove noise_scale from kwargs so it doesn't get passed to the parent class
kwargs.pop("noise_scale", None)
super().__init__(*args, **kwargs)
def add_noise(
self,
original_samples: torch.Tensor,
noise: torch.Tensor,
timesteps: torch.Tensor,
) -> torch.Tensor:
t_01 = (timesteps / 1000).to(original_samples.device)
scaled_noise = noise * self.noise_scale
noisy_model_input = (1.0 - t_01) * original_samples + t_01 * scaled_noise
return noisy_model_input
def add_special_tokens(tokenizer):
"""Attach the special-token shortcuts that the pipeline relies on."""
tokenizer.boi_token = "<|boi_token|>"
tokenizer.bor_token = "<|bor_token|>"
tokenizer.eor_token = "<|eor_token|>"
tokenizer.bot_token = "<|bot_token|>"
tokenizer.tms_token = "<|tms_token|>"
def get_tokenizer(processor):
from transformers import PreTrainedTokenizerBase
if isinstance(processor, PreTrainedTokenizerBase):
return processor
return processor.tokenizer
class FakeConfig:
pass
class FakeTextEncoder(torch.nn.Module):
def __init__(self, scaling_factor=1.0):
super().__init__()
self._dtype = torch.float32
self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.config = FakeConfig()
self.config.scaling_factor = scaling_factor
@property
def dtype(self):
return self._dtype
@dtype.setter
def dtype(self, value):
self._dtype = value
@property
def device(self):
return self._device
@device.setter
def device(self, value):
self._device = value
# mimic to from torch
def to(self, *args, **kwargs):
# pull out dtype and device if they exist
if "dtype" in kwargs:
self._dtype = kwargs["dtype"]
if "device" in kwargs:
self._device = kwargs["device"]
return super().to(*args, **kwargs)
class HidreamO1Model(BaseModel):
arch = "hidream_o1"
def __init__(
self,
device,
model_config: ModelConfig,
dtype="bf16",
custom_pipeline=None,
noise_scheduler=None,
**kwargs,
):
super().__init__(
device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
)
self.is_flow_matching = True
self.is_transformer = True
self.target_lora_modules = ["Qwen3VLForConditionalGeneration"]
self.noise_scale = self.model_config.model_kwargs.get(
"noise_scale", DEFAULT_NOISE_SCALE
)
self.noise_scale_inference = self.model_config.model_kwargs.get(
"noise_scale_inference", self.noise_scale
)
print(f"Using noise scale: {self.noise_scale}")
global _GLOBAL_NOISE_SCALE
_GLOBAL_NOISE_SCALE = self.noise_scale
self.is_comfy_weight = False # save as single file if true
# static method to get the noise scheduler
@staticmethod
def get_train_scheduler():
return HidreamO1FlowmatchScheduler(
**scheduler_config, noise_scale=_GLOBAL_NOISE_SCALE
)
def get_bucket_divisibility(self):
return 32 # patch size
def load_model(self):
dtype = self.torch_dtype
self.print_and_status_update("Loading HidreamO1 model")
model_path = self.model_config.name_or_path
self.print_and_status_update("Loading transformer")
try:
processor = AutoProcessor.from_pretrained(model_path)
except Exception as e:
print(
f"Failed to load processor from model path {model_path}, trying original path. Error: {e}"
)
processor_path = self.model_config.extras_name_or_path
if processor_path.endswith(".safetensors"):
processor_path = "HiDream-ai/HiDream-O1-Image"
processor = AutoProcessor.from_pretrained(processor_path)
tokenizer = get_tokenizer(processor)
add_special_tokens(tokenizer)
if model_path.endswith(".safetensors"):
self.is_comfy_weight = True
self.print_and_status_update(
"Model is in safetensors format, loading with safetensors"
)
state_dict = load_file(model_path)
for key, value in state_dict.items():
state_dict[key] = value.to(dtype=dtype)
# comfy ui is missing the lm head. It isnt used, but our model needs it for now
state_dict["lm_head.weight"] = torch.zeros(
151936, 4096, dtype=torch.bfloat16, device="cpu"
)
# transformer.load_state_dict(state_dict, assign=True)
transformer = Qwen3VLForConditionalGeneration.from_pretrained(
None,
config=Qwen3VLConfig(**model_config),
state_dict=state_dict,
torch_dtype=self.torch_dtype,
)
del state_dict # free memory
else:
transformer = Qwen3VLForConditionalGeneration.from_pretrained(
model_path,
torch_dtype=self.torch_dtype,
)
flush()
if not self.model_config.low_vram:
transformer.to(self.device_torch)
if self.model_config.quantize:
self.print_and_status_update("Quantizing Transformer")
quantize_model(self, transformer)
flush()
if (
self.model_config.layer_offloading
and self.model_config.layer_offloading_transformer_percent > 0
):
MemoryManager.attach(
transformer,
self.device_torch,
offload_percent=self.model_config.layer_offloading_transformer_percent,
ignore_modules=[],
)
flush()
# move over to device now if low vram
if self.model_config.low_vram:
transformer.to(self.device_torch)
# fake ones so the trainer doesnt break
vae = FakeVAE().to(self.device_torch, dtype=dtype)
text_encoder = FakeTextEncoder().to(self.device_torch, dtype=dtype)
self.noise_scheduler = HidreamO1Model.get_train_scheduler()
self.print_and_status_update("Making pipe")
kwargs = {}
pipe: HiDreamO1Pipeline = HiDreamO1Pipeline(
scheduler=self.noise_scheduler,
processor=processor,
model=None,
**kwargs,
)
pipe.model = transformer
self.print_and_status_update("Preparing Model")
flush()
# save it to the model class
self.vae = vae
self.text_encoder = text_encoder
self.tokenizer = processor
self.model = pipe.model
self.pipeline = pipe
self.print_and_status_update("Model Loaded")
def get_generation_pipeline(self):
scheduler = HidreamO1Model.get_train_scheduler()
pipe: HiDreamO1Pipeline = HiDreamO1Pipeline(
scheduler=scheduler,
processor=self.tokenizer,
model=None,
)
pipe.model = self.transformer
return pipe
def encode_images(self, image_list: torch.Tensor, device=None, dtype=None):
if self.vae.device == torch.device("cpu"):
self.vae.to(self.device_torch)
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
# not needed since there is not a latent space
return image_list.to(device, dtype=dtype)
def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
if self.vae.device == torch.device("cpu"):
self.vae.to(self.device_torch)
if device is None:
device = self.vae_device_torch
if dtype is None:
dtype = self.vae_torch_dtype
# not needed since there is not a latent space
return latents.to(device, dtype=dtype)
def generate_single_image(
self,
pipeline: HiDreamO1Pipeline,
gen_config: GenerateImageConfig,
conditional_embeds: AdvancedPromptEmbeds,
unconditional_embeds: AdvancedPromptEmbeds,
generator: torch.Generator,
extra: dict,
):
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
sc = self.get_bucket_divisibility()
gen_config.width = int(gen_config.width // sc * sc)
gen_config.height = int(gen_config.height // sc * sc)
img = pipeline(
# prompt=gen_config.prompt,
prompt_input_ids=conditional_embeds.text_embeds[0],
# negative_prompt=gen_config.negative_prompt,
negative_prompt_input_ids=unconditional_embeds.text_embeds[0],
height=gen_config.height,
width=gen_config.width,
num_inference_steps=gen_config.num_inference_steps,
guidance_scale=gen_config.guidance_scale,
generator=generator,
noise_scale=self.noise_scale_inference,
**extra,
).images[0]
return img
def get_noise_prediction(
self,
latent_model_input: torch.Tensor,
timestep: torch.Tensor, # 0 to 1000 scale
text_embeddings: AdvancedPromptEmbeds,
batch: "DataLoaderBatchDTO",
**kwargs,
):
import einops
from .src.hidream_o1.pipeline import PATCH_SIZE, T_EPS
if self.model.device == torch.device("cpu"):
self.model.to(self.device_torch)
device = self.device_torch
in_dtype = latent_model_input.dtype
bs, _, h_pix, w_pix = latent_model_input.shape
h_patches = h_pix // PATCH_SIZE
w_patches = w_pix // PATCH_SIZE
# (B, C, H, W) -> (B, H/p * W/p, C * p * p)
z = einops.rearrange(
latent_model_input,
"B C (H p1) (W p2) -> B (H W) (C p1 p2)",
p1=PATCH_SIZE,
p2=PATCH_SIZE,
).to(device)
model_config = self.model.config
pad_token_id = getattr(model_config, "pad_token_id", 0) or 0
with torch.no_grad():
# Build per-sample conditioning, then left-pad the text portion so
# the boi/tms + vision-token suffix stays at the end of the
# sequence (the t2i layout assumes vision tokens are at the tail).
per_sample = []
for b in range(bs):
tokens = text_embeddings.text_embeds[b]
if tokens.dim() == 1:
tokens = tokens.unsqueeze(0)
per_sample.append(
self.pipeline.build_conditioning_sample(
tokens.to(device),
h_pix,
w_pix,
)
)
max_seq_len = max(s["input_ids"].shape[-1] for s in per_sample)
ids_l, pos_l, tt_l, vm_l, mask_l = [], [], [], [], []
for s in per_sample:
ids = s["input_ids"].to(device)
pos = s["position_ids"].to(device)
tt = s["token_types"].to(device)
vm = s["vinput_mask"].to(device)
seq_len = ids.shape[-1]
pad_len = max_seq_len - seq_len
if pad_len > 0:
ids = torch.cat(
[
torch.full(
(1, pad_len),
pad_token_id,
dtype=ids.dtype,
device=device,
),
ids,
],
dim=-1,
)
pos = torch.cat(
[
torch.ones((3, 1, pad_len), dtype=pos.dtype, device=device),
pos,
],
dim=-1,
)
tt = torch.cat(
[
torch.zeros((1, pad_len), dtype=tt.dtype, device=device),
tt,
],
dim=-1,
)
vm = torch.cat(
[
torch.zeros((1, pad_len), dtype=vm.dtype, device=device),
vm,
],
dim=-1,
)
mask = torch.cat(
[
torch.zeros((1, pad_len), dtype=torch.long, device=device),
torch.ones((1, seq_len), dtype=torch.long, device=device),
],
dim=-1,
)
else:
mask = torch.ones((1, seq_len), dtype=torch.long, device=device)
ids_l.append(ids)
pos_l.append(pos)
tt_l.append(tt)
vm_l.append(vm)
mask_l.append(mask)
input_ids = torch.cat(ids_l, dim=0)
position_ids = torch.cat(pos_l, dim=1) # (3, B, S)
token_types = torch.cat(tt_l, dim=0)
vinput_mask = torch.cat(vm_l, dim=0)
attention_mask = torch.cat(mask_l, dim=0)
# Model wants timestep as denoising progress in (0, 1) where 1=clean.
t_pixeldit = (1.0 - timestep.float() / 1000.0).to(device)
outputs = self.model(
input_ids=input_ids,
position_ids=position_ids,
attention_mask=attention_mask if bs > 1 else None,
vinputs=z,
timestep=t_pixeldit.reshape(-1),
token_types=token_types,
use_flash_attn=True,
)
x_pred = outputs.x_pred # (B, S, C*p*p) over the full padded sequence
# Pull the vision-token positions only.
vision_pred = torch.stack(
[x_pred[b][vinput_mask[b].bool()] for b in range(bs)],
dim=0,
) # (B, image_len, C*p*p)
x0_pred = einops.rearrange(
vision_pred,
"B (H W) (C p1 p2) -> B C (H p1) (W p2)",
H=h_patches,
W=w_patches,
p1=PATCH_SIZE,
p2=PATCH_SIZE,
)
# Model emits an x0-prediction; convert to flow-matching velocity
# (x_1 - x_0) so it matches the loss target from get_loss_target.
sigma = (timestep.float() / 1000.0).clamp_min(T_EPS).to(device)
while sigma.dim() < latent_model_input.dim():
sigma = sigma.unsqueeze(-1)
pred = (latent_model_input.float().to(device) - x0_pred.float()) / sigma
return pred.to(in_dtype)
def get_prompt_embeds(self, prompt: list) -> AdvancedPromptEmbeds:
if not isinstance(prompt, list):
prompt = [prompt]
# empty, we cannot use them with this omni model anyway, but will break trainer if they do not exist
token_list = [self.pipeline.encode_prompt(p) for p in prompt]
pe = AdvancedPromptEmbeds(text_embeds=token_list)
pe._frozen_dtype_keys = ["text_embeds"]
return pe
def get_model_has_grad(self):
return False
def get_te_has_grad(self):
return False
def save_model(self, output_path, meta, save_dtype):
transformer: Qwen3VLForConditionalGeneration = unwrap_model(self.model)
if self.is_comfy_weight:
sd = transformer.state_dict()
save_dict = {}
for key, value in sd.items():
if "lm_head.weight" in key:
continue # comfy checkpoint doesnt have the lm head, so skip it
save_dict[key] = value.clone().to("cpu", dtype=save_dtype)
if not output_path.endswith(".safetensors"):
output_path += ".safetensors"
meta = get_meta_for_safetensors(meta, name=self.arch)
save_file(save_dict, output_path, metadata=meta)
else:
transformer.save_pretrained(
save_directory=output_path,
safe_serialization=True,
)
# save processor
self.tokenizer.save_pretrained(output_path)
meta_path = os.path.join(output_path, "aitk_meta.yaml")
with open(meta_path, "w") as f:
yaml.dump(meta, f)
def get_loss_target(self, *args, **kwargs):
noise = kwargs.get("noise")
batch = kwargs.get("batch")
noise_scale = self.noise_scale
return (noise * noise_scale - batch.latents).detach()
def get_base_model_version(self):
return self.arch
def get_transformer_block_names(self) -> Optional[List[str]]:
return ["layers"]
def convert_lora_weights_before_save(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("transformer.", "diffusion_model.")
new_key = new_key.replace(".model.", ".")
new_sd[new_key] = value
return new_sd
def convert_lora_weights_before_load(self, state_dict):
new_sd = {}
for key, value in state_dict.items():
new_key = key.replace("diffusion_model.", "transformer.model.")
# to load legacy keys
new_key = new_key.replace("transformer.model.model.", "transformer.model.")
new_sd[new_key] = value
return new_sd

0
extensions_built_in/diffusion_models/hidream/src/hidream_o1/__init__.py Normal file
View File

52
extensions_built_in/diffusion_models/hidream/src/hidream_o1/model_config.py Normal file
View File

@@ -0,0 +1,52 @@
model_config = {
"architectures": ["Qwen3VLForConditionalGeneration"],
"image_token_id": 151655,
"model_type": "qwen3_vl",
"text_config": {
"attention_bias": False,
"attention_dropout": 0.0,
"bos_token_id": 151643,
"dtype": "bfloat16",
"eos_token_id": 151645,
"head_dim": 128,
"hidden_act": "silu",
"hidden_size": 4096,
"initializer_range": 0.02,
"intermediate_size": 12288,
"max_position_embeddings": 262144,
"model_type": "qwen3_vl_text",
"num_attention_heads": 32,
"num_hidden_layers": 36,
"num_key_value_heads": 8,
"rms_norm_eps": 1e-06,
"rope_scaling": {
"mrope_interleaved": True,
"mrope_section": [24, 20, 20],
"rope_type": "default",
},
"rope_theta": 5000000,
"use_cache": True,
"vocab_size": 151936,
},
"tie_word_embeddings": False,
"transformers_version": "4.57.0.dev0",
"video_token_id": 151656,
"vision_config": {
"deepstack_visual_indexes": [8, 16, 24],
"depth": 27,
"hidden_act": "gelu_pytorch_tanh",
"hidden_size": 1152,
"in_channels": 3,
"initializer_range": 0.02,
"intermediate_size": 4304,
"model_type": "qwen3_vl",
"num_heads": 16,
"num_position_embeddings": 2304,
"out_hidden_size": 4096,
"patch_size": 16,
"spatial_merge_size": 2,
"temporal_patch_size": 2,
},
"vision_end_token_id": 151653,
"vision_start_token_id": 151652,
}

455
extensions_built_in/diffusion_models/hidream/src/hidream_o1/pipeline.py Normal file
View File

@@ -0,0 +1,455 @@
from typing import List, Optional, Union
import einops
import numpy as np
import torch
from PIL import Image
import torchvision.transforms.v2 as transforms
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from diffusers.utils import BaseOutput
from dataclasses import dataclass
TIMESTEP_TOKEN_NUM = 1
DEFAULT_NOISE_SCALE = 8.0
T_EPS = 0.001
PATCH_SIZE = 32
TENSOR_TRANSFORM = transforms.Compose(
[
transforms.ToImage(),
transforms.ToDtype(torch.float32, scale=True),
transforms.Normalize([0.5], [0.5]),
]
)
def round_to_patch(dim: int, patch: int = PATCH_SIZE) -> int:
return max(patch, int(dim // patch * patch))
def _get_rope_index_t2i(
spatial_merge_size: int,
image_token_id: int,
video_token_id: int,
vision_start_token_id: int,
input_ids: torch.LongTensor,
image_grid_thw: torch.LongTensor,
skip_vision_start_token: List[int],
fix_point: int = 4096,
):
"""Compute mrope position ids for the t2i case used by HiDream-O1."""
attention_mask = torch.ones_like(input_ids)
position_ids = torch.ones(
3,
input_ids.shape[0],
input_ids.shape[1],
dtype=input_ids.dtype,
device=input_ids.device,
)
for i, ids_row in enumerate(input_ids):
ids_row = ids_row[attention_mask[i] == 1]
vision_start_indices = torch.argwhere(ids_row == vision_start_token_id).squeeze(
1
)
vision_tokens = ids_row[vision_start_indices + 1]
image_nums = (vision_tokens == image_token_id).sum().item()
video_nums = (vision_tokens == video_token_id).sum().item()
input_tokens = ids_row.tolist()
llm_pos_ids_list = []
st = 0
image_index = 0
video_index = 0
remain_images, remain_videos = image_nums, video_nums
local_fix_point = fix_point
for _ in range(image_nums + video_nums):
ed_image = (
input_tokens.index(image_token_id, st)
if (image_token_id in input_tokens and remain_images > 0)
else len(input_tokens) + 1
)
ed_video = (
input_tokens.index(video_token_id, st)
if (video_token_id in input_tokens and remain_videos > 0)
else len(input_tokens) + 1
)
if ed_image < ed_video:
t, h, w = image_grid_thw[image_index].tolist()
image_index += 1
remain_images -= 1
ed = ed_image
else:
t, h, w = image_grid_thw[video_index].tolist()
video_index += 1
remain_videos -= 1
ed = ed_video
llm_grid_t = t
llm_grid_h = h // spatial_merge_size
llm_grid_w = w // spatial_merge_size
text_len = ed - st - skip_vision_start_token[image_index - 1]
text_len = max(0, text_len)
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
)
t_index = (
torch.arange(llm_grid_t)
.view(-1, 1)
.expand(-1, llm_grid_h * llm_grid_w)
.flatten()
)
h_index = (
torch.arange(llm_grid_h)
.view(1, -1, 1)
.expand(llm_grid_t, -1, llm_grid_w)
.flatten()
)
w_index = (
torch.arange(llm_grid_w)
.view(1, 1, -1)
.expand(llm_grid_t, llm_grid_h, -1)
.flatten()
)
if skip_vision_start_token[image_index - 1]:
if local_fix_point > 0:
local_fix_point = local_fix_point - st_idx
llm_pos_ids_list.append(
torch.stack([t_index, h_index, w_index]) + local_fix_point + st_idx
)
local_fix_point = 0
else:
llm_pos_ids_list.append(
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
)
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
if st < len(input_tokens):
st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
text_len = len(input_tokens) - st
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
)
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
position_ids.device
)
return position_ids
def _build_t2i_sample_from_input_ids(
input_ids: torch.Tensor,
height: int,
width: int,
model_config,
attention_mask: Optional[torch.Tensor] = None,
):
"""Build the full conditioning sample (position_ids/token_types/vinput_mask)
around an already-tokenized prompt."""
image_token_id = model_config.image_token_id
video_token_id = model_config.video_token_id
vision_start_token_id = model_config.vision_start_token_id
image_len = (height // PATCH_SIZE) * (width // PATCH_SIZE)
if input_ids.dim() == 1:
input_ids = input_ids.unsqueeze(0)
image_grid_thw = torch.tensor(
[1, height // PATCH_SIZE, width // PATCH_SIZE], dtype=torch.int64
).unsqueeze(0)
vision_tokens = (
torch.zeros((1, image_len), dtype=input_ids.dtype, device=input_ids.device)
+ image_token_id
)
vision_tokens[0, 0] = vision_start_token_id
input_ids_pad = torch.cat([input_ids, vision_tokens], dim=-1)
position_ids = _get_rope_index_t2i(
spatial_merge_size=1,
image_token_id=image_token_id,
video_token_id=video_token_id,
vision_start_token_id=vision_start_token_id,
input_ids=input_ids_pad,
image_grid_thw=image_grid_thw,
skip_vision_start_token=[1],
)
txt_seq_len = input_ids.shape[-1]
all_seq_len = position_ids.shape[-1]
token_types = torch.zeros((1, all_seq_len), dtype=input_ids.dtype)
bgn = txt_seq_len - TIMESTEP_TOKEN_NUM
token_types[0, bgn : bgn + image_len + TIMESTEP_TOKEN_NUM] = 1
token_types[0, txt_seq_len - TIMESTEP_TOKEN_NUM : txt_seq_len] = 3
vinput_mask = token_types == 1
token_types_bin = (token_types > 0).to(token_types.dtype)
sample = {
"input_ids": input_ids,
"position_ids": position_ids,
"token_types": token_types_bin,
"vinput_mask": vinput_mask,
}
if attention_mask is not None:
if attention_mask.dim() == 1:
attention_mask = attention_mask.unsqueeze(0)
sample["attention_mask"] = attention_mask
return sample
@dataclass
class HiDreamO1PipelineOutput(BaseOutput):
images: List[Image.Image]
class HiDreamO1Pipeline(DiffusionPipeline):
"""
Diffusers-style inference pipeline for HiDream-O1 (base model).
HiDream-O1 is a unified text/vision/diffusion model with no VAE — the
transformer directly predicts image patches in pixel space. This pipeline
keeps only the components needed for text-to-image inference.
"""
model_cpu_offload_seq = "model"
def __init__(
self,
model,
processor,
scheduler: FlowMatchEulerDiscreteScheduler,
):
super().__init__()
self.register_modules(model=model, processor=processor, scheduler=scheduler)
@property
def tokenizer(self):
return (
self.processor.tokenizer
if hasattr(self.processor, "tokenizer")
else self.processor
)
def _snap_resolution(self, width: int, height: int):
w, h = round_to_patch(width), round_to_patch(height)
if (w, h) != (width, height):
print(f"[hidream-o1] Resolution rounded from {width}x{height} to {w}x{h}")
return w, h
def build_conditioning_sample(
self,
input_ids: torch.Tensor,
height: int,
width: int,
attention_mask: Optional[torch.Tensor] = None,
):
"""Build the per-sample conditioning dict (input_ids, position_ids,
token_types, vinput_mask) around already-tokenized text. Useful when
a training loop needs to batch samples manually."""
return _build_t2i_sample_from_input_ids(
input_ids,
height,
width,
self.model.config,
attention_mask=attention_mask,
)
def encode_prompt(self, prompt: str) -> torch.Tensor:
"""Apply the chat template + boi/tms suffix and tokenize.
Returns input_ids of shape (1, seq_len). Use these to precompute and
pass back into __call__ via `prompt_input_ids` / `negative_prompt_input_ids`."""
tokenizer = self.tokenizer
boi_token = getattr(tokenizer, "boi_token", "<|boi_token|>")
tms_token = getattr(tokenizer, "tms_token", "<|tms_token|>")
messages = [{"role": "user", "content": prompt}]
template_caption = (
self.processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
+ boi_token
+ tms_token * TIMESTEP_TOKEN_NUM
)
return tokenizer.encode(
template_caption, return_tensors="pt", add_special_tokens=False
)
@torch.no_grad()
def __call__(
self,
prompt: Optional[Union[str, List[str]]] = None,
negative_prompt: Optional[Union[str, List[str]]] = " ",
prompt_input_ids: Optional[torch.Tensor] = None,
negative_prompt_input_ids: Optional[torch.Tensor] = None,
prompt_attention_mask: Optional[torch.Tensor] = None,
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
height: int = 1440,
width: int = 2560,
num_inference_steps: int = 50,
guidance_scale: float = 5.0,
shift: float = 3.0,
generator: Optional[torch.Generator] = None,
seed: Optional[int] = None,
noise_scale: float = None,
output_type: str = "pil",
return_dict: bool = True,
):
if noise_scale is None:
noise_scale = DEFAULT_NOISE_SCALE
if prompt is None and prompt_input_ids is None:
raise ValueError("Provide either `prompt` or `prompt_input_ids`.")
def _unwrap_str(x):
if isinstance(x, list):
if len(x) != 1:
raise ValueError(
"HiDreamO1Pipeline currently supports batch size 1."
)
return x[0]
return x
prompt = _unwrap_str(prompt)
negative_prompt = _unwrap_str(negative_prompt)
device = self._execution_device
dtype = torch.bfloat16
model_config = self.model.config
width, height = self._snap_resolution(width, height)
h_patches = height // PATCH_SIZE
w_patches = width // PATCH_SIZE
do_cfg = guidance_scale > 1.0
if prompt_input_ids is None:
prompt_input_ids = self.encode_prompt(prompt)
if do_cfg and negative_prompt_input_ids is None:
if negative_prompt is None:
negative_prompt = " "
negative_prompt_input_ids = self.encode_prompt(negative_prompt)
cond_sample = _build_t2i_sample_from_input_ids(
prompt_input_ids,
height,
width,
model_config,
attention_mask=prompt_attention_mask,
)
uncond_sample = (
_build_t2i_sample_from_input_ids(
negative_prompt_input_ids,
height,
width,
model_config,
attention_mask=negative_prompt_attention_mask,
)
if do_cfg
else None
)
def _to_device(s):
return {
k: (v.to(device) if torch.is_tensor(v) else v) for k, v in s.items()
}
cond_sample = _to_device(cond_sample)
if uncond_sample is not None:
uncond_sample = _to_device(uncond_sample)
if generator is None:
if seed is None:
seed = 0
generator = torch.Generator(device="cpu").manual_seed(seed + 1)
torch.manual_seed(seed + 1)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed + 1)
noise = noise_scale * torch.randn(
(1, 3, height, width), generator=generator
).to(device, dtype)
z = einops.rearrange(
noise,
"B C (H p1) (W p2) -> B (H W) (C p1 p2)",
p1=PATCH_SIZE,
p2=PATCH_SIZE,
)
if shift is not None and hasattr(self.scheduler, "set_shift"):
self.scheduler.set_shift(shift)
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
def _forward_once(sample, z_in, t_pixeldit):
with torch.autocast(device.type, dtype=dtype):
kwargs = {
"input_ids": sample["input_ids"],
"position_ids": sample["position_ids"],
"vinputs": z_in,
"timestep": t_pixeldit.reshape(-1).to(device),
"token_types": sample["token_types"],
"use_flash_attn": True,
}
if "attention_mask" in sample:
kwargs["attention_mask"] = sample["attention_mask"]
outputs = self.model(**kwargs)
x_pred = outputs.x_pred
return x_pred[0, sample["vinput_mask"][0]].unsqueeze(0)
for step_t in self.progress_bar(timesteps):
t_pixeldit = 1.0 - step_t.float() / 1000.0
sigma = (step_t.float() / 1000.0).to(dtype=torch.float32).clamp_min(T_EPS)
x_pred_cond = _forward_once(cond_sample, z.clone(), t_pixeldit)
v_cond = (x_pred_cond.float() - z.float()) / sigma
if do_cfg:
x_pred_uncond = _forward_once(uncond_sample, z.clone(), t_pixeldit)
v_uncond = (x_pred_uncond.float() - z.float()) / sigma
v_guided = v_uncond + guidance_scale * (v_cond - v_uncond)
else:
v_guided = v_cond
model_output = -v_guided
z = self.scheduler.step(
model_output.float(),
step_t.to(dtype=torch.float32),
z.float(),
return_dict=False,
)[0].to(dtype)
img = (z + 1) / 2
img = einops.rearrange(
img.cpu().float(),
"B (H W) (C p1 p2) -> B C (H p1) (W p2)",
H=h_patches,
W=w_patches,
p1=PATCH_SIZE,
p2=PATCH_SIZE,
)
if output_type == "pt":
images = img.clamp(0, 1)
elif output_type == "np":
images = np.clip(img.numpy().transpose(0, 2, 3, 1), 0, 1)
else:
arr = np.round(
np.clip(img[0].numpy().transpose(1, 2, 0) * 255, 0, 255)
).astype(np.uint8)
images = [Image.fromarray(arr).convert("RGB")]
if not return_dict:
return (images,)
return HiDreamO1PipelineOutput(images=images)

2465
extensions_built_in/diffusion_models/hidream/src/hidream_o1/qwen3_vl_transformers.py Normal file
View File

File diff suppressed because it is too large Load Diff

12
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -916,7 +916,17 @@ class SDTrainer(BaseSDTrainProcess):
loss = loss + norm_std_loss
return loss + additional_loss
loss = loss + additional_loss
if self.train_config.max_loss_debug and self.train_config.max_loss is not None:
if loss.item() > self.train_config.max_loss:
print_acc(f"Loss {loss.item()} is greater than max loss {self.train_config.max_loss}. Clipping to max loss.")
print_acc(f"timesteps: {timesteps}")
if self.train_config.max_loss is not None:
loss = torch.clamp(loss, max=self.train_config.max_loss)
return loss
def preprocess_batch(self, batch: 'DataLoaderBatchDTO'):
return batch

47
toolkit/advanced_prompt_embeds.py
View File

@@ -1,4 +1,5 @@
import os
import torch
from safetensors.torch import load_file, save_file
@@ -30,11 +31,20 @@ class AdvancedPromptEmbeds:
def __init__(self, **kwargs):
self._store = {}
self._frozen_dtype_keys = []
for key, value in kwargs.items():
if not isinstance(value, list):
value = [value]
self._store[key] = value
@property
def frozen_dtype_keys(self):
return self._frozen_dtype_keys
@frozen_dtype_keys.setter
def frozen_dtype_keys(self, keys):
self._frozen_dtype_keys = list(keys) if keys else []
def __getattr__(self, name):
if name.startswith("_"):
raise AttributeError(name)
@@ -46,10 +56,14 @@ class AdvancedPromptEmbeds:
def __setattr__(self, name, value):
if name.startswith("_"):
super().__setattr__(name, value)
else:
if not isinstance(value, list):
value = [value]
self._store[name] = value
return
cls_attr = getattr(type(self), name, None)
if isinstance(cls_attr, property):
super().__setattr__(name, value)
return
if not isinstance(value, list):
value = [value]
self._store[name] = value
def set(self, key, value):
if not isinstance(value, list):
@@ -77,25 +91,38 @@ class AdvancedPromptEmbeds:
return key in self._store
def to(self, *args, **kwargs):
frozen = set(self._frozen_dtype_keys)
if frozen:
no_dtype_args = [a for a in args if not isinstance(a, torch.dtype)]
no_dtype_kwargs = {k: v for k, v in kwargs.items() if k != "dtype"}
new_pe = AdvancedPromptEmbeds()
new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys)
for key, value in self._store.items():
new_pe._store[key] = [v.to(*args, **kwargs) for v in value]
if key in frozen:
new_pe._store[key] = [
v.to(*no_dtype_args, **no_dtype_kwargs) for v in value
]
else:
new_pe._store[key] = [v.to(*args, **kwargs) for v in value]
return new_pe
def detach(self):
new_pe = AdvancedPromptEmbeds()
new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys)
for key, value in self._store.items():
new_pe._store[key] = [v.detach() for v in value]
return new_pe
def clone(self):
new_pe = AdvancedPromptEmbeds()
new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys)
for key, value in self._store.items():
new_pe._store[key] = [v.clone() for v in value]
return new_pe
def expand_to_batch(self, batch_size):
new_pe = AdvancedPromptEmbeds()
new_pe._frozen_dtype_keys = list(self._frozen_dtype_keys)
for key, value in self._store.items():
if len(value) == 1:
new_pe._store[key] = value * batch_size
@@ -137,12 +164,18 @@ class AdvancedPromptEmbeds:
cls, prompt_embeds: list["AdvancedPromptEmbeds"], padding_side: str = "right"
):
embeds = {}
frozen = []
for pe in prompt_embeds:
for key in pe.keys():
if key not in embeds:
embeds[key] = []
embeds[key].extend(pe[key])
return cls(**embeds)
for k in pe.frozen_dtype_keys:
if k not in frozen:
frozen.append(k)
out = cls(**embeds)
out.frozen_dtype_keys = frozen
return out
@classmethod
def split_prompt_embeds(cls, concatenated: "AdvancedPromptEmbeds", num_parts=None):
@@ -150,6 +183,8 @@ class AdvancedPromptEmbeds:
# use length of first item as num_parts
num_parts = len(concatenated[concatenated.keys()[0]])
split_embeds = [cls() for _ in range(num_parts)]
for pe in split_embeds:
pe.frozen_dtype_keys = list(concatenated.frozen_dtype_keys)
for key in concatenated.keys():
values = concatenated[key]
if len(values) != num_parts:

5
toolkit/config_modules.py
View File

@@ -577,6 +577,11 @@ class TrainConfig:
self.do_blank_stabilization = kwargs.get('do_blank_stabilization', False)
self.audio_loss_multiplier = kwargs.get("audio_loss_multiplier", 1.0)
# will throw detailed error when it goes over
self.max_loss_debug: bool = kwargs.get("max_loss_debug", False)
# will clip the loss to this amount to prevent wild outliers
self.max_loss: Optional[float] = kwargs.get("max_loss", None)
ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21']

32
ui/src/app/jobs/new/SimpleJob.tsx
View File

@@ -380,20 +380,24 @@ export default function SimpleJob({
}}
options={transformerQuantizationOptions}
/>
<SelectInput
label="Text Encoder"
value={jobConfig.config.process[0].model.quantize_te ? jobConfig.config.process[0].model.qtype_te : ''}
onChange={value => {
if (value === '') {
setJobConfig(false, 'config.process[0].model.quantize_te');
value = defaultQtype;
} else {
setJobConfig(true, 'config.process[0].model.quantize_te');
{!disableSections.includes('model.quantize_te') && (
<SelectInput
label="Text Encoder"
value={
jobConfig.config.process[0].model.quantize_te ? jobConfig.config.process[0].model.qtype_te : ''
}
setJobConfig(value, 'config.process[0].model.qtype_te');
}}
options={quantizationOptions}
/>
onChange={value => {
if (value === '') {
setJobConfig(false, 'config.process[0].model.quantize_te');
value = defaultQtype;
} else {
setJobConfig(true, 'config.process[0].model.quantize_te');
}
setJobConfig(value, 'config.process[0].model.qtype_te');
}}
options={quantizationOptions}
/>
)}
</Card>
)}
{modelArch?.additionalSections?.includes('model.multistage') && (
@@ -1068,7 +1072,7 @@ export default function SimpleJob({
<div className="grid grid-cols-2 gap-2">
{[
[256, 512, 768, 1024],
[1280, 1328, 1536],
[1280, 1328, 1536, 2048],
].map(resGroup => (
<div key={resGroup[0]} className="space-y-2">
{resGroup.map(res => (

35
ui/src/app/jobs/new/options.ts
View File

@@ -6,6 +6,7 @@ type Control = 'depth' | 'line' | 'pose' | 'inpaint';
type DisableableSections =
| 'model.quantize'
| 'model.quantize_te'
| 'train.timestep_type'
| 'network.conv'
| 'trigger_word'
@@ -922,6 +923,40 @@ export const modelArchs: ModelArch[] = [
disableSections: ['network.conv'],
additionalSections: ['model.low_vram'],
},
{
name: 'hidream_o1',
label: 'HiDream-O1',
group: 'image',
defaults: {
'config.process[0].model.name_or_path': ['HiDream-ai/HiDream-O1-Image', defaultNameOrPath],
'config.process[0].model.quantize': [true, false],
'config.process[0].model.quantize_te': [false, false],
'config.process[0].train.timestep_type': ['linear', 'sigmoid'],
'config.process[0].network.conv': [undefined, 16],
'config.process[0].network.conv_alpha': [undefined, 16],
'config.process[0].train.max_loss': [1.0, undefined],
'config.process[0].network.network_kwargs.ignore_if_contains': [['lm_head','patch_embed', 'visual'], []],
'config.process[0].network.transformer_only': [false, undefined],
'config.process[0].sample.width': [2048, 1024],
'config.process[0].sample.height': [2048, 1024],
'config.process[0].model.model_kwargs': [
{
noise_scale_inference: 8.0,
noise_scale: 8.0,
},
{},
],
},
disableSections: [
'network.conv',
'model.quantize_te',
'train.unload_text_encoder',
],
additionalSections: [
'model.low_vram',
'model.layer_offloading',
],
},
].sort((a, b) => {
// Sort by label, case-insensitive
return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });

2
ui/src/types.ts
View File

@@ -71,6 +71,7 @@ export interface NetworkConfig {
network_kwargs: {
ignore_if_contains: string[];
};
transformer_only?: boolean;
}
export interface SaveConfig {
@@ -149,6 +150,7 @@ export interface TrainConfig {
do_differential_guidance?: boolean;
differential_guidance_scale?: number;
audio_loss_multiplier?: number;
max_loss?: number | null;
}
export interface QuantizeKwargsConfig {

2
version.py
View File

@@ -1 +1 @@
VERSION = "0.9.8"
VERSION = "0.9.10"