diff --git a/modules/dataLoader/IdeogramBaseDataLoader.py b/modules/dataLoader/IdeogramBaseDataLoader.py new file mode 100644 index 000000000..28de23e40 --- /dev/null +++ b/modules/dataLoader/IdeogramBaseDataLoader.py @@ -0,0 +1,159 @@ +import os + +from modules.dataLoader.BaseDataLoader import BaseDataLoader +from modules.dataLoader.mixin.DataLoaderText2ImageMixin import DataLoaderText2ImageMixin +from modules.model.IdeogramModel import PROMPT_MAX_LENGTH, IdeogramModel +from modules.modelSetup.BaseIdeogramSetup import BaseIdeogramSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.TrainProgress import TrainProgress + +from mgds.pipelineModules.DecodeTokens import DecodeTokens +from mgds.pipelineModules.DecodeVAE import DecodeVAE +from mgds.pipelineModules.EncodeIdeogramText import EncodeIdeogramText +from mgds.pipelineModules.EncodeVAE import EncodeVAE +from mgds.pipelineModules.PadMaskedTokens import PadMaskedTokens +from mgds.pipelineModules.PruneMaskedTokens import PruneMaskedTokens +from mgds.pipelineModules.RescaleImageChannels import RescaleImageChannels +from mgds.pipelineModules.SampleVAEDistribution import SampleVAEDistribution +from mgds.pipelineModules.SaveImage import SaveImage +from mgds.pipelineModules.SaveText import SaveText +from mgds.pipelineModules.ScaleImage import ScaleImage +from mgds.pipelineModules.Tokenize import Tokenize + + +@factory.register(BaseDataLoader, ModelType.IDEOGRAM_4) +class IdeogramBaseDataLoader( + BaseDataLoader, + DataLoaderText2ImageMixin, +): + def _preparation_modules(self, config: TrainConfig, model: IdeogramModel): + rescale_image = RescaleImageChannels(image_in_name='image', image_out_name='image', in_range_min=0, in_range_max=1, out_range_min=-1, out_range_max=1) + encode_image = EncodeVAE(in_name='image', out_name='latent_image_distribution', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + image_sample = SampleVAEDistribution(in_name='latent_image_distribution', out_name='latent_image', mode='mean') + downscale_mask = ScaleImage(in_name='mask', out_name='latent_mask', factor=0.125) + + # Ideogram wraps each prompt in the Qwen chat template (mirrors Ideogram4Pipeline.encode_prompt / encode_text). + # EncodeIdeogramText uses a plain arange for RoPE positions, and encode_text left-aligns by reading [:, :n], + # both of which require the real tokens to come first — relies on the tokenizer's default right-padding. + tokenize_prompt = Tokenize( + in_name='prompt', tokens_out_name='tokens', mask_out_name='tokens_mask', + tokenizer=model.tokenizer, max_token_length=PROMPT_MAX_LENGTH, + apply_chat_template=lambda prompt: [{"role": "user", "content": [{"type": "text", "text": prompt}]}], + apply_chat_template_kwargs={'add_generation_prompt': True}, + ) + encode_prompt = EncodeIdeogramText( + tokens_name='tokens', tokens_attention_mask_in_name='tokens_mask', + hidden_state_out_name='text_encoder_hidden_state', tokens_attention_mask_out_name='tokens_mask', + text_encoder=model.text_encoder, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype(), + ) + prune_masked_tokens = PruneMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state') + + modules = [rescale_image, encode_image, image_sample] + if config.masked_training or config.model_type.has_mask_input(): + modules.append(downscale_mask) + modules += [tokenize_prompt, encode_prompt] + + if config.latent_caching: + modules.append(prune_masked_tokens) + + return modules + + def _cache_modules(self, config: TrainConfig, model: IdeogramModel, model_setup: BaseIdeogramSetup): + image_split_names = ['latent_image', 'original_resolution', 'crop_offset'] + + if config.masked_training or config.model_type.has_mask_input(): + image_split_names.append('latent_mask') + + image_aggregate_names = ['crop_resolution', 'image_path'] + + text_split_names = [] + + sort_names = image_aggregate_names + image_split_names + [ + 'prompt', 'tokens', 'tokens_mask', 'text_encoder_hidden_state', + 'concept' + ] + + text_split_names += ['tokens', 'tokens_mask', 'text_encoder_hidden_state'] + + return self._cache_modules_from_names( + model, model_setup, + image_split_names=image_split_names, + image_aggregate_names=image_aggregate_names, + text_split_names=text_split_names, + sort_names=sort_names, + config=config, + text_caching=True, + ) + + def _output_modules(self, config: TrainConfig, model: IdeogramModel, model_setup: BaseIdeogramSetup): + pad_masked_tokens = PadMaskedTokens(tokens_name='tokens', tokens_mask_name='tokens_mask', hidden_state_name='text_encoder_hidden_state', max_length=PROMPT_MAX_LENGTH) + + output_names = [ + 'image_path', 'latent_image', + 'prompt', + 'tokens', + 'tokens_mask', + 'original_resolution', 'crop_resolution', 'crop_offset', + ] + + if config.masked_training or config.model_type.has_mask_input(): + output_names.append('latent_mask') + + output_names.append('text_encoder_hidden_state') + + output_module_list = self._output_modules_from_out_names( + model, model_setup, + output_names=output_names, + config=config, + use_conditioning_image=False, + vae=model.vae, + autocast_context=[model.autocast_context], + train_dtype=model.train_dtype, + ) + + if config.latent_caching: + output_module_list = [pad_masked_tokens] + output_module_list + + return output_module_list + + def _debug_modules(self, config: TrainConfig, model: IdeogramModel): + debug_dir = os.path.join(config.debug_dir, "dataloader") + + def before_save_fun(): + model.vae_to(self.train_device) + + decode_image = DecodeVAE(in_name='latent_image', out_name='decoded_image', vae=model.vae, autocast_contexts=[model.autocast_context], dtype=model.train_dtype.torch_dtype()) + upscale_mask = ScaleImage(in_name='latent_mask', out_name='decoded_mask', factor=8) + decode_prompt = DecodeTokens(in_name='tokens', out_name='decoded_prompt', tokenizer=model.tokenizer) + save_image = SaveImage(image_in_name='decoded_image', original_path_in_name='image_path', path=debug_dir, in_range_min=-1, in_range_max=1, before_save_fun=before_save_fun) + save_mask = SaveImage(image_in_name='decoded_mask', original_path_in_name='image_path', path=debug_dir, in_range_min=0, in_range_max=1, before_save_fun=before_save_fun) + save_prompt = SaveText(text_in_name='decoded_prompt', original_path_in_name='image_path', path=debug_dir, before_save_fun=before_save_fun) + + modules = [] + + modules.append(decode_image) + modules.append(save_image) + + if config.masked_training or config.model_type.has_mask_input(): + modules.append(upscale_mask) + modules.append(save_mask) + + modules.append(decode_prompt) + modules.append(save_prompt) + + return modules + + def _create_dataset( + self, + config: TrainConfig, + model: IdeogramModel, + model_setup: BaseIdeogramSetup, + train_progress: TrainProgress, + is_validation: bool = False, + ): + return DataLoaderText2ImageMixin._create_dataset(self, + config, model, model_setup, train_progress, is_validation, + aspect_bucketing_quantization=64, + ) diff --git a/modules/model/IdeogramModel.py b/modules/model/IdeogramModel.py new file mode 100644 index 000000000..d388615f8 --- /dev/null +++ b/modules/model/IdeogramModel.py @@ -0,0 +1,276 @@ +import math +from contextlib import nullcontext +from random import Random + +PROMPT_MAX_LENGTH = 2048 + +from modules.model.BaseModel import BaseModel +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util.enum.ModelType import ModelType +from modules.util.LayerOffloadConductor import LayerOffloadConductor + +import torch +from torch import Tensor + +from diffusers import ( + AutoencoderKLFlux2, + DiffusionPipeline, + FlowMatchEulerDiscreteScheduler, + Ideogram4Pipeline, + Ideogram4Transformer2DModel, +) +from diffusers.pipelines.ideogram4.pipeline_ideogram4 import _resolution_aware_mu +from transformers import Qwen2TokenizerFast, Qwen3VLModel + + +class IdeogramModel(BaseModel): + # base model data + tokenizer: Qwen2TokenizerFast | None + noise_scheduler: FlowMatchEulerDiscreteScheduler | None + text_encoder: Qwen3VLModel | None + vae: AutoencoderKLFlux2 | None + transformer: Ideogram4Transformer2DModel | None + # second, frozen DiT for the asymmetric (dual-network) classifier-free guidance — only used at sampling + unconditional_transformer: Ideogram4Transformer2DModel | None + + # autocast context + text_encoder_autocast_context: torch.autocast | nullcontext + + text_encoder_offload_conductor: LayerOffloadConductor | None + transformer_offload_conductor: LayerOffloadConductor | None + unconditional_transformer_offload_conductor: LayerOffloadConductor | None + + transformer_lora: LoRAModuleWrapper | None + lora_state_dict: dict | None + + def __init__( + self, + model_type: ModelType, + ): + super().__init__( + model_type=model_type, + ) + + self.tokenizer = None + self.noise_scheduler = None + self.text_encoder = None + self.vae = None + self.transformer = None + self.unconditional_transformer = None + + self.text_encoder_autocast_context = nullcontext() + + self.text_encoder_offload_conductor = None + self.transformer_offload_conductor = None + self.unconditional_transformer_offload_conductor = None + + self.transformer_lora = None + self.lora_state_dict = None + + def adapters(self) -> list[LoRAModuleWrapper]: + # only the conditional transformer is trainable; the unconditional transformer never sees the concept + return [a for a in [ + self.transformer_lora, + ] if a is not None] + + def fusion_groups(self) -> list | None: + # Ideogram4 fuses q/k/v into one qkv Linear per block; everything else in the transformer -- including + # the output projection (to_out.0 -> o, see diffusers_to_original) -- already matches the original + # checkpoint's naming. + return [ + ("layers.{i}", ["attention.to_q", "attention.to_k", "attention.to_v"], "attention.qkv", "attention.qkv"), + ] + + def diffusers_to_original(self) -> list | None: + # Ideogram4's native (ComfyUI/original) checkpoint is identical to diffusers except for one rename: + # to_out.0 -> o on the (already fused, see fusion_groups above) attention module. Everything else -- + # feed_forward, norms, adaln_modulation, final_layer, t_embedding, etc. -- keeps its diffusers name, so + # a single catch-all identity pattern covers the rest. + return [ + ("layers.{i}.attention.to_out.0", "layers.{i}.attention.o"), + ("{path}", "{path}"), + ] + + def vae_to(self, device: torch.device): + self.vae.to(device=device) + + def text_encoder_to(self, device: torch.device): + if self.text_encoder_offload_conductor is not None: + self.text_encoder_offload_conductor.to(device) + else: + self.text_encoder.to(device=device) + + def transformer_to(self, device: torch.device): + if self.transformer_offload_conductor is not None: + self.transformer_offload_conductor.to(device) + else: + self.transformer.to(device=device) + + if self.transformer_lora is not None: + self.transformer_lora.to(device) + + def unconditional_transformer_to(self, device: torch.device): + if self.unconditional_transformer is not None: + if self.unconditional_transformer_offload_conductor is not None: + self.unconditional_transformer_offload_conductor.to(device) + else: + self.unconditional_transformer.to(device=device) + + def to(self, device: torch.device): + self.vae_to(device) + self.text_encoder_to(device) + self.transformer_to(device) + self.unconditional_transformer_to(device) + + def eval(self): + self.vae.eval() + self.text_encoder.eval() + self.transformer.eval() + if self.unconditional_transformer is not None: + self.unconditional_transformer.eval() + + def create_pipeline(self) -> DiffusionPipeline: + return Ideogram4Pipeline( + transformer=self.transformer, + unconditional_transformer=self.unconditional_transformer, + vae=self.vae, + text_encoder=self.text_encoder, + tokenizer=self.tokenizer, + scheduler=self.noise_scheduler, + prompt_enhancer_head=None, + ) + + def encode_text( + self, + train_device: torch.device, + batch_size: int = 1, + rand: Random | None = None, + text: str | list[str] | None = None, + tokens: Tensor | None = None, + tokens_mask: Tensor | None = None, + text_encoder_dropout_probability: float | None = None, + text_encoder_output: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + if tokens is None and text is not None: + if isinstance(text, str): + text = [text] + + # Ideogram wraps each prompt in a chat template, like Z-Image. Source: Ideogram4Pipeline.encode_prompt. + for i, prompt_item in enumerate(text): + messages = [{"role": "user", "content": [{"type": "text", "text": prompt_item}]}] + text[i] = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + + tokenizer_output = self.tokenizer( + text, + padding='max_length', + max_length=PROMPT_MAX_LENGTH, + truncation=True, + return_tensors='pt', + add_special_tokens=False, + ) + tokens = tokenizer_output.input_ids.to(self.text_encoder.device) + tokens_mask = tokenizer_output.attention_mask.to(self.text_encoder.device) + + if text_encoder_output is None: + with self.text_encoder_autocast_context: + # Mirrors Ideogram4Pipeline.encode_prompt: text-only MRoPE shares the linear token position across + # all 3 axes, so a plain arange is the position_ids for the real tokens. + position_ids = torch.arange(tokens.shape[1], device=tokens.device).unsqueeze(0).expand(tokens.shape[0], -1) + selected = Ideogram4Pipeline._get_text_encoder_hidden_states(self.text_encoder, tokens, tokens_mask, position_ids) + + # Interleave by hidden dim (NOT torch.cat): stack -> (L, B, T, H), permute -> (B, T, H, L), reshape -> + # (B, T, H*L). This grouped-by-hidden-dim order is what the transformer's llm_cond_proj expects. + text_encoder_output = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(tokens.shape[0], tokens.shape[1], -1) + + if text_encoder_dropout_probability is not None and text_encoder_dropout_probability > 0.0: + raise NotImplementedError # needs empty-caption conditioning, not zero-out + + # Run unconditionally in both the fresh and cached paths: zero out padding positions. + text_encoder_output = text_encoder_output * tokens_mask.to(text_encoder_output.dtype).unsqueeze(-1) + + text_lengths = tokens_mask.sum(dim=1).long() + + # Left-align each sample's real tokens to the END of the text block, matching the [left-pad][text][image] + # layout that prepare_packed_ids (Ideogram4Pipeline._prepare_ids) builds. The transformer masks padding via + # segment_ids, so the padding's sequence position is otherwise irrelevant; aligning to the ids' convention here + # lets every caller pack the features directly, with no second alignment pass (the single-prompt sampler has + # offset 0, so this is a no-op there). Real tokens are front-aligned coming in (right-padded tokenizer output). + max_text_tokens = int(text_lengths.max().item()) + aligned_output = text_encoder_output.new_zeros( + text_encoder_output.shape[0], max_text_tokens, text_encoder_output.shape[-1], + ) + for b, num_text in enumerate(text_lengths.tolist()): + aligned_output[b, max_text_tokens - num_text:] = text_encoder_output[b, :num_text] + return aligned_output, text_lengths + + @staticmethod + def patchify_latents(latents: torch.Tensor) -> torch.Tensor: + # (B, C, H, W) VAE latents -> packed (B, L, C * 4) sequence with L = (H // 2) * (W // 2). + # Inverse of the unpatchify in Ideogram4Pipeline.__call__'s decode tail. + b, c, h, w = latents.shape + latents = latents.view(b, c, h // 2, 2, w // 2, 2) + latents = latents.permute(0, 2, 4, 3, 5, 1) + return latents.reshape(b, (h // 2) * (w // 2), c * 4) + + @staticmethod + def prepare_packed_ids( + text_lengths: list[int] | Tensor, + grid_h: int, + grid_w: int, + max_text_tokens: int, + device: torch.device, + ) -> tuple[Tensor, Tensor, Tensor]: + # Ideogram4Pipeline._prepare_ids takes a plain list[int]; encode_text returns text_lengths as a Tensor. + if isinstance(text_lengths, Tensor): + text_lengths = text_lengths.tolist() + return Ideogram4Pipeline._prepare_ids(text_lengths, grid_h, grid_w, max_text_tokens, device) + + @staticmethod + def pack_llm_features(text_features: Tensor, num_image_tokens: int) -> Tensor: + # Assemble the packed [text][image] conditioning: the image positions carry zeroed text features (they are + # conditioned via position_ids/segment_ids/indicator, not features). Mirrors Ideogram4Pipeline.encode_prompt. + # The caller supplies text_features already aligned to the target layout (front-aligned for the single-prompt + # sampler, left-padded for the batched training predict). + image_feature_padding = torch.zeros( + text_features.shape[0], num_image_tokens, text_features.shape[-1], + dtype=text_features.dtype, device=text_features.device, + ) + return torch.cat([text_features, image_feature_padding], dim=1) + + @staticmethod + def unpatchify_latents(latents: torch.Tensor, grid_h: int, grid_w: int) -> torch.Tensor: + # packed (B, L, C * 4) -> (B, C, H, W). Mirrors Ideogram4Pipeline.__call__'s decode unpatchify. + b, _l, cp = latents.shape + c = cp // 4 + latents = latents.view(b, grid_h, grid_w, 2, 2, c) + latents = latents.permute(0, 5, 1, 3, 2, 4) + return latents.reshape(b, c, grid_h * 2, grid_w * 2) + + def scale_latents(self, latents: Tensor) -> Tensor: + # Operates on packed (B, L, C * 4); vae.bn stats are per packed-channel, so (C*4,) broadcasts on the last dim. + # Same batch-norm de/normalization as Flux2Model.scale_latents (same AutoencoderKLFlux2), duplicated here only + # because that copy uses the unpacked (1, -1, 1, 1) broadcast; matches Ideogram4Pipeline.__call__'s bn scaling. + mean = self.vae.bn.running_mean.view(1, 1, -1).to(latents.device, latents.dtype) + std = torch.sqrt( + self.vae.bn.running_var.view(1, 1, -1) + self.vae.config.batch_norm_eps + ).to(latents.device, latents.dtype) + return (latents - mean) / std + + def unscale_latents(self, latents: Tensor) -> Tensor: + mean = self.vae.bn.running_mean.view(1, 1, -1).to(latents.device, latents.dtype) + std = torch.sqrt( + self.vae.bn.running_var.view(1, 1, -1) + self.vae.config.batch_norm_eps + ).to(latents.device, latents.dtype) + return latents * std + mean + + def calculate_timestep_shift(self, latent_height: int, latent_width: int) -> float: + # Ideogram shifts the flow-matching schedule by a resolution-aware mu (pipeline _resolution_aware_mu, relative to + # a 512x512 base). OneTrainer's flow-matching timestep sampling applies a multiplicative shift, which is exp(mu) + # (same mu->shift mapping as Ernie). _resolution_aware_mu only depends on the pixel-count ratio, so convert the + # latent dims back to pixels (VAE scale factor 8). + mu = _resolution_aware_mu(height=latent_height * 8, width=latent_width * 8, base_mu=0.0) + return math.exp(mu) diff --git a/modules/modelLoader/IdeogramModelLoader.py b/modules/modelLoader/IdeogramModelLoader.py new file mode 100644 index 000000000..5a91d51c2 --- /dev/null +++ b/modules/modelLoader/IdeogramModelLoader.py @@ -0,0 +1,192 @@ +import os +import traceback + +from modules.model.IdeogramModel import IdeogramModel +from modules.modelLoader.GenericFineTuneModelLoader import make_fine_tune_model_loader +from modules.modelLoader.GenericLoRAModelLoader import make_lora_model_loader +from modules.modelLoader.mixin.HFModelLoaderMixin import HFModelLoaderMixin +from modules.modelLoader.mixin.LoRALoaderMixin import LoRALoaderMixin +from modules.util.config.TrainConfig import QuantizationConfig +from modules.util.enum.ModelType import ModelType +from modules.util.ModelNames import ModelNames +from modules.util.ModelWeightDtypes import ModelWeightDtypes + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + Ideogram4Transformer2DModel, +) +from transformers import AutoTokenizer, Qwen3VLModel + + +class IdeogramModelLoader( + HFModelLoaderMixin, +): + def __init__(self): + super().__init__() + + def __load_internal( + self, + model: IdeogramModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + include_unconditional_transformer: bool, + quantization: QuantizationConfig, + ): + if os.path.isfile(os.path.join(base_model_name, "meta.json")): + self.__load_diffusers(model, model_type, weight_dtypes, base_model_name, include_unconditional_transformer, quantization) + else: + raise Exception("not an internal model") + + def __load_diffusers( + self, + model: IdeogramModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + include_unconditional_transformer: bool, + quantization: QuantizationConfig, + ): + transformer = self._load_diffusers_sub_module( + Ideogram4Transformer2DModel, + weight_dtypes.transformer, + weight_dtypes.train_dtype, + base_model_name, + "transformer", + quantization, + ) + # the unconditional transformer is frozen and only used for the negative branch of the dual-network CFG at + # sampling, so it has its own weight dtype independent of the trainable transformer's. It is optional: if not + # loaded, only cfg_scale<=1 sampling is possible. + if include_unconditional_transformer: + unconditional_transformer = self._load_diffusers_sub_module( + Ideogram4Transformer2DModel, + weight_dtypes.unconditional_transformer, + weight_dtypes.train_dtype, + base_model_name, + "unconditional_transformer", + quantization, + ) + else: + unconditional_transformer = None + + text_encoder = self._load_transformers_sub_module( + Qwen3VLModel, + weight_dtypes.text_encoder, + weight_dtypes.fallback_train_dtype, + base_model_name, + "text_encoder", + ) + + tokenizer = AutoTokenizer.from_pretrained( + base_model_name, + subfolder="tokenizer", + ) + + noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + base_model_name, + subfolder="scheduler", + ) + + vae = self._load_diffusers_sub_module( + AutoencoderKLFlux2, + weight_dtypes.vae, + weight_dtypes.train_dtype, + base_model_name, + "vae", + ) + + model.model_type = model_type + model.tokenizer = tokenizer + model.noise_scheduler = noise_scheduler + model.text_encoder = text_encoder + model.vae = vae + model.transformer = transformer + model.unconditional_transformer = unconditional_transformer + + def __load_safetensors( + self, + model: IdeogramModel, + model_type: ModelType, + weight_dtypes: ModelWeightDtypes, + base_model_name: str, + quantization: QuantizationConfig, + ): + raise NotImplementedError( + "Loading single-file safetensors for Ideogram is not supported. Use the diffusers model instead." + ) + + def load( + self, + model: IdeogramModel, + model_type: ModelType, + model_names: ModelNames, + weight_dtypes: ModelWeightDtypes, + quantization: QuantizationConfig, + ): + stacktraces = [] + + try: + self.__load_internal( + model, model_type, weight_dtypes, model_names.base_model, + model_names.include_unconditional_transformer, quantization, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + try: + self.__load_diffusers( + model, model_type, weight_dtypes, model_names.base_model, + model_names.include_unconditional_transformer, quantization, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + try: + self.__load_safetensors( + model, model_type, weight_dtypes, model_names.base_model, quantization, + ) + return + except Exception: + stacktraces.append(traceback.format_exc()) + + for stacktrace in stacktraces: + print(stacktrace) + raise Exception("could not load model: " + model_names.base_model) + + +class IdeogramLoRALoader( + LoRALoaderMixin +): + def __init__(self): + super().__init__() + + def load( + self, + model: IdeogramModel, + model_names: ModelNames, + ): + return self._load(model, model_names) + + +IdeogramLoRAModelLoader = make_lora_model_loader( + model_spec_map={ + ModelType.IDEOGRAM_4: "resources/sd_model_spec/ideogram_4-lora.json", + }, + model_class=IdeogramModel, + model_loader_class=IdeogramModelLoader, + lora_loader_class=IdeogramLoRALoader, + embedding_loader_class=None, +) + +IdeogramFineTuneModelLoader = make_fine_tune_model_loader( + model_spec_map={ + ModelType.IDEOGRAM_4: "resources/sd_model_spec/ideogram_4.json", + }, + model_class=IdeogramModel, + model_loader_class=IdeogramModelLoader, + embedding_loader_class=None, +) diff --git a/modules/modelSampler/IdeogramSampler.py b/modules/modelSampler/IdeogramSampler.py new file mode 100644 index 000000000..0a5dd420b --- /dev/null +++ b/modules/modelSampler/IdeogramSampler.py @@ -0,0 +1,244 @@ +import copy +from collections.abc import Callable + +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSampler.BaseModelSampler import BaseModelSampler, ModelSamplerOutput +from modules.util import factory +from modules.util.config.SampleConfig import SampleConfig +from modules.util.enum.AudioFormat import AudioFormat +from modules.util.enum.FileType import FileType +from modules.util.enum.ImageFormat import ImageFormat +from modules.util.enum.ModelType import ModelType +from modules.util.enum.NoiseScheduler import NoiseScheduler +from modules.util.enum.VideoFormat import VideoFormat +from modules.util.torch_util import torch_gc + +import torch + +from diffusers.pipelines.ideogram4.pipeline_ideogram4 import _logit_normal_sigmas, _resolution_aware_mu + +import numpy as np +from PIL import Image as PILImage +from tqdm import tqdm + + +@factory.register(BaseModelSampler, ModelType.IDEOGRAM_4) +class IdeogramSampler(BaseModelSampler): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + model: IdeogramModel, + model_type: ModelType, + ): + super().__init__(train_device, temp_device) + + self.model = model + self.model_type = model_type + self.pipeline = model.create_pipeline() + + @torch.no_grad() + def __sample_base( + self, + prompt: str, + negative_prompt: str, + height: int, + width: int, + seed: int, + random_seed: bool, + diffusion_steps: int, + cfg_scale: float, + noise_scheduler: NoiseScheduler, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ) -> ModelSamplerOutput: + with self.model.autocast_context: + generator = torch.Generator(device=self.train_device) + if random_seed: + generator.seed() + else: + generator.manual_seed(seed) + + noise_scheduler = copy.deepcopy(self.model.noise_scheduler) + vae = self.pipeline.vae + transformer = self.pipeline.transformer + dtype = self.model.train_dtype.torch_dtype() + + # Ideogram uses asymmetric (dual-network) CFG: the negative branch is normally the unconditional_transformer + # run on the image tokens with zeroed text features, NOT a negative-prompt encode. negative_prompt is + # unused. If the unconditional transformer was not loaded, fall back to encoding an empty ("") prompt and + # running it through the conditional transformer instead, like standard CFG. + use_cfg = cfg_scale > 1.0 + use_unconditional_transformer = use_cfg and self.model.unconditional_transformer is not None + use_empty_prompt_negative = use_cfg and self.model.unconditional_transformer is None + + vae_scale_factor = 8 + patch_size = 2 + latent_dim = transformer.config.in_channels + grid_h = height // (vae_scale_factor * patch_size) + grid_w = width // (vae_scale_factor * patch_size) + num_image_tokens = grid_h * grid_w + + # build the packed [text][image] conditioning for a single text encode. Padding positions are masked + # out by segment_ids/indicator, so packing to the actual text length matches the 2048-pad pipeline. + def pack_conditioning(text_features: torch.Tensor, text_lengths: torch.Tensor) -> tuple: + max_text_tokens = text_features.shape[1] + position_ids, segment_ids, indicator = self.model.prepare_packed_ids( + text_lengths, grid_h, grid_w, max_text_tokens, self.train_device, + ) + llm_features = self.model.pack_llm_features(text_features, num_image_tokens).to(dtype) + text_z_padding = torch.zeros( + text_features.shape[0], max_text_tokens, latent_dim, dtype=dtype, device=self.train_device, + ) + return max_text_tokens, position_ids, segment_ids, indicator, llm_features, text_z_padding + + # encode text (conditional branch, and the empty-prompt negative branch if needed) + self.model.text_encoder_to(self.train_device) + text_features, text_lengths = self.model.encode_text( + train_device=self.train_device, + text=prompt, + ) + max_text_tokens, position_ids, segment_ids, indicator, llm_features, text_z_padding = pack_conditioning( + text_features, text_lengths, + ) + + if use_empty_prompt_negative: + neg_text_features, neg_text_lengths = self.model.encode_text( + train_device=self.train_device, + text="", + ) + ( + max_neg_text_tokens, neg_position_ids, neg_segment_ids, neg_indicator, neg_llm_features, + neg_text_z_padding, + ) = pack_conditioning(neg_text_features, neg_text_lengths) + del neg_text_features + self.model.text_encoder_to(self.temp_device) + torch_gc() + + if use_unconditional_transformer: + # unconditional (image-only) branch: zeroed text features over the image-region slices of the layout + neg_position_ids = position_ids[:, max_text_tokens:] + neg_segment_ids = segment_ids[:, max_text_tokens:] + neg_indicator = indicator[:, max_text_tokens:] + neg_llm_features = torch.zeros( + text_features.shape[0], num_image_tokens, text_features.shape[-1], + dtype=dtype, device=self.train_device, + ) + + # packed (B, num_image_tokens, latent_dim) noise + latent_image = torch.randn( + size=(text_features.shape[0], num_image_tokens, latent_dim), + generator=generator, device=self.train_device, dtype=torch.float32, + ) + + # free before the denoising loop; closes the gap on the OOM observed in llm_cond_norm's fp32 variance + # upcast of neg_llm_features + del text_features + + # resolution-aware logit-normal Euler schedule (pipeline overrides the scheduler's default sigmas) + schedule_mu = _resolution_aware_mu(height=height, width=width, base_mu=0.0) + sigmas = _logit_normal_sigmas(diffusion_steps, schedule_mu, std=1.5, device=self.train_device) + noise_scheduler.set_timesteps(sigmas=sigmas.tolist(), device=self.train_device) + timesteps = noise_scheduler.timesteps + num_train_timesteps = noise_scheduler.config.num_train_timesteps + + self.model.transformer_to(self.train_device) + if use_unconditional_transformer: + self.model.unconditional_transformer_to(self.train_device) + + for i, timestep in enumerate(tqdm(timesteps, desc="sampling")): + # scheduler stores num_train_timesteps-scaled timesteps; convert back to model time (0=noise, 1=data) + t_model = (1.0 - timestep.float() / num_train_timesteps).expand(latent_image.shape[0]) + + pos_z = torch.cat([text_z_padding, latent_image.to(dtype)], dim=1) + pos_out = transformer( + hidden_states=pos_z, + timestep=t_model, + encoder_hidden_states=llm_features, + position_ids=position_ids, + segment_ids=segment_ids, + indicator=indicator, + return_dict=False, + )[0] + pos_v = pos_out[:, max_text_tokens:].float() + + if use_unconditional_transformer: + neg_v = self.model.unconditional_transformer( + hidden_states=latent_image.to(dtype), + timestep=t_model, + encoder_hidden_states=neg_llm_features, + position_ids=neg_position_ids, + segment_ids=neg_segment_ids, + indicator=neg_indicator, + return_dict=False, + )[0].float() + elif use_empty_prompt_negative: + neg_z = torch.cat([neg_text_z_padding, latent_image.to(dtype)], dim=1) + neg_out = transformer( + hidden_states=neg_z, + timestep=t_model, + encoder_hidden_states=neg_llm_features, + position_ids=neg_position_ids, + segment_ids=neg_segment_ids, + indicator=neg_indicator, + return_dict=False, + )[0] + neg_v = neg_out[:, max_neg_text_tokens:].float() + + v = neg_v + cfg_scale * (pos_v - neg_v) if use_cfg else pos_v + + latent_image = noise_scheduler.step(-v, timestep, latent_image, return_dict=False)[0] + + on_update_progress(i + 1, len(timesteps)) + + self.model.transformer_to(self.temp_device) + self.model.unconditional_transformer_to(self.temp_device) + torch_gc() + self.model.vae_to(self.train_device) + + # bn-denormalize the packed latents and unpatchify back to (B, C, H, W) before VAE decode + latents = self.model.unscale_latents(latent_image) + latents = self.model.unpatchify_latents(latents, grid_h, grid_w) + + image = vae.decode(latents.to(vae.dtype), return_dict=False)[0] + # no VaeImageProcessor — match the pipeline's manual postprocess + image = (image.clamp(-1, 1) + 1) / 2 + image = image.cpu().permute(0, 2, 3, 1).float().numpy() + image = [PILImage.fromarray((img * 255).astype(np.uint8)) for img in image] + + self.model.vae_to(self.temp_device) + torch_gc() + + return ModelSamplerOutput( + file_type=FileType.IMAGE, + data=image[0], + ) + + def sample( + self, + sample_config: SampleConfig, + destination: str, + image_format: ImageFormat | None = None, + video_format: VideoFormat | None = None, + audio_format: AudioFormat | None = None, + on_sample: Callable[[ModelSamplerOutput], None] = lambda _: None, + on_update_progress: Callable[[int, int], None] = lambda _, __: None, + ): + sampler_output = self.__sample_base( + prompt=sample_config.prompt, + negative_prompt=sample_config.negative_prompt, + height=self.quantize_resolution(sample_config.height, 64), + width=self.quantize_resolution(sample_config.width, 64), + seed=sample_config.seed, + random_seed=sample_config.random_seed, + diffusion_steps=sample_config.diffusion_steps, + cfg_scale=sample_config.cfg_scale, + noise_scheduler=sample_config.noise_scheduler, + on_update_progress=on_update_progress, + ) + + self.save_sampler_output( + sampler_output, destination, + image_format, video_format, audio_format, + ) + + on_sample(sampler_output) diff --git a/modules/modelSaver/IdeogramFineTuneModelSaver.py b/modules/modelSaver/IdeogramFineTuneModelSaver.py new file mode 100644 index 000000000..29f38ea01 --- /dev/null +++ b/modules/modelSaver/IdeogramFineTuneModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSaver.GenericFineTuneModelSaver import make_fine_tune_model_saver +from modules.modelSaver.ideogram.IdeogramModelSaver import IdeogramModelSaver +from modules.util.enum.ModelType import ModelType + +IdeogramFineTuneModelSaver = make_fine_tune_model_saver( + ModelType.IDEOGRAM_4, + model_class=IdeogramModel, + model_saver_class=IdeogramModelSaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/IdeogramLoRAModelSaver.py b/modules/modelSaver/IdeogramLoRAModelSaver.py new file mode 100644 index 000000000..e18e25027 --- /dev/null +++ b/modules/modelSaver/IdeogramLoRAModelSaver.py @@ -0,0 +1,11 @@ +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSaver.GenericLoRAModelSaver import make_lora_model_saver +from modules.modelSaver.ideogram.IdeogramLoRASaver import IdeogramLoRASaver +from modules.util.enum.ModelType import ModelType + +IdeogramLoRAModelSaver = make_lora_model_saver( + ModelType.IDEOGRAM_4, + model_class=IdeogramModel, + lora_saver_class=IdeogramLoRASaver, + embedding_saver_class=None, +) diff --git a/modules/modelSaver/ideogram/IdeogramLoRASaver.py b/modules/modelSaver/ideogram/IdeogramLoRASaver.py new file mode 100644 index 000000000..8d9a712d6 --- /dev/null +++ b/modules/modelSaver/ideogram/IdeogramLoRASaver.py @@ -0,0 +1,23 @@ +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSaver.mixin.LoRASaverMixin import LoRASaverMixin + +from torch import Tensor + + +class IdeogramLoRASaver( + LoRASaverMixin, +): + def __init__(self): + super().__init__() + + def _get_state_dict( + self, + model: IdeogramModel, + ) -> dict[str, Tensor]: + # only the conditional transformer is trained; the unconditional transformer never sees the concept + state_dict = {} + if model.transformer_lora is not None: + state_dict |= model.transformer_lora.state_dict() + if model.lora_state_dict is not None: + state_dict |= model.lora_state_dict + return state_dict diff --git a/modules/modelSaver/ideogram/IdeogramModelSaver.py b/modules/modelSaver/ideogram/IdeogramModelSaver.py new file mode 100644 index 000000000..ec625e5bd --- /dev/null +++ b/modules/modelSaver/ideogram/IdeogramModelSaver.py @@ -0,0 +1,72 @@ +import os.path +from pathlib import Path + +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSaver.mixin.DtypeModelSaverMixin import DtypeModelSaverMixin +from modules.util.enum.ModelFormat import ModelFormat + +import torch + +from safetensors.torch import save_file + + +class IdeogramModelSaver( + DtypeModelSaverMixin, +): + def __init__(self): + super().__init__() + + def __save_diffusers( + self, + model: IdeogramModel, + destination: str, + dtype: torch.dtype | None, + ): + pipeline = model.create_pipeline() + pipeline.to("cpu") + save_pipeline = self._copy_pipeline_to_dtype(pipeline, dtype, pipeline.tokenizer) + + os.makedirs(Path(destination).absolute(), exist_ok=True) + save_pipeline.save_pretrained(destination) + + if dtype is not None: + del save_pipeline + + def __save_safetensors( + self, + model: IdeogramModel, + destination: str, + dtype: torch.dtype | None, + ): + state_dict = model.transformer.state_dict() + + save_state_dict = self._convert_state_dict_dtype(state_dict, dtype) + self._convert_state_dict_to_contiguous(save_state_dict) + + os.makedirs(Path(destination).parent.absolute(), exist_ok=True) + + save_file(save_state_dict, destination, self._create_safetensors_header(model, save_state_dict)) + + def __save_internal( + self, + model: IdeogramModel, + destination: str, + ): + self.__save_diffusers(model, destination, None) + + def save( + self, + model: IdeogramModel, + output_model_format: ModelFormat, + output_model_destination: str, + dtype: torch.dtype | None, + ): + match output_model_format: + case ModelFormat.DIFFUSERS: + self.__save_diffusers(model, output_model_destination, dtype) + case ModelFormat.ORIGINAL_TRANSFORMER: + self.__save_safetensors(model, output_model_destination, dtype) + case ModelFormat.INTERNAL: + self.__save_internal(model, output_model_destination) + case _: + raise NotImplementedError(f"Unsupported output format: {output_model_format}") diff --git a/modules/modelSetup/BaseIdeogramSetup.py b/modules/modelSetup/BaseIdeogramSetup.py new file mode 100644 index 000000000..279c09388 --- /dev/null +++ b/modules/modelSetup/BaseIdeogramSetup.py @@ -0,0 +1,219 @@ +from abc import ABCMeta +from random import Random + +import modules.util.multi_gpu_util as multi +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.modelSetup.mixin.ModelSetupDebugMixin import ModelSetupDebugMixin +from modules.modelSetup.mixin.ModelSetupDiffusionLossMixin import ModelSetupDiffusionLossMixin +from modules.modelSetup.mixin.ModelSetupEmbeddingMixin import ModelSetupEmbeddingMixin +from modules.modelSetup.mixin.ModelSetupFlowMatchingMixin import ModelSetupFlowMatchingMixin +from modules.modelSetup.mixin.ModelSetupNoiseMixin import ModelSetupNoiseMixin +from modules.util.checkpointing_util import ( + enable_checkpointing_for_ideogram_transformer, + enable_checkpointing_for_qwen3vl_encoder_layers, +) +from modules.util.config.TrainConfig import TrainConfig +from modules.util.dtype_util import create_autocast_context, disable_fp16_autocast_context +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.quantization_util import quantize_layers +from modules.util.torch_util import torch_gc +from modules.util.TrainProgress import TrainProgress + +import torch +from torch import Tensor + + +class BaseIdeogramSetup( + BaseModelSetup, + ModelSetupDiffusionLossMixin, + ModelSetupDebugMixin, + ModelSetupNoiseMixin, + ModelSetupFlowMatchingMixin, + ModelSetupEmbeddingMixin, + metaclass=ABCMeta +): + LAYER_PRESETS = { + "attn-mlp": ["attention", "feed_forward"], + "attn-only": ["attention"], + "blocks": ["layers"], + "full": [], + } + + def setup_optimizations( + self, + model: IdeogramModel, + config: TrainConfig, + ): + # Only the conditional transformer is trained, so gradient checkpointing applies there. + model.transformer_offload_conductor = \ + enable_checkpointing_for_ideogram_transformer(model.transformer, config, config.transformer) + + # The unconditional transformer is frozen, but it still benefits from layer offloading + # since both transformers need to fit in VRAM during sampling. It is optional, so may be unloaded. + if model.unconditional_transformer is not None: + model.unconditional_transformer_offload_conductor = \ + enable_checkpointing_for_ideogram_transformer(model.unconditional_transformer, config, config.unconditional_transformer) + + model.text_encoder_offload_conductor = enable_checkpointing_for_qwen3vl_encoder_layers(model.text_encoder, config, config.text_encoder) + + model.autocast_context, model.train_dtype = create_autocast_context(self.train_device, config.train_dtype, [ + config.weight_dtypes().transformer, + config.weight_dtypes().text_encoder, + config.weight_dtypes().vae, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], config.enable_autocast_cache) + + model.text_encoder_autocast_context, model.text_encoder_train_dtype = \ + disable_fp16_autocast_context( + self.train_device, + config.train_dtype, + config.fallback_train_dtype, + [ + config.weight_dtypes().text_encoder, + config.weight_dtypes().lora if config.training_method == TrainingMethod.LORA else None, + ], + config.enable_autocast_cache, + ) + + quantize_layers(model.text_encoder, self.train_device, model.text_encoder_train_dtype, config) + quantize_layers(model.vae, self.train_device, model.train_dtype, config) + quantize_layers(model.transformer, self.train_device, model.train_dtype, config) + quantize_layers(model.unconditional_transformer, self.train_device, model.train_dtype, config) + + self._set_attention_backend(model.transformer, config.attention_mechanism, mask=False) + if model.unconditional_transformer is not None: + self._set_attention_backend(model.unconditional_transformer, config.attention_mechanism, mask=False) + + def predict( + self, + model: IdeogramModel, + batch: dict, + config: TrainConfig, + train_progress: TrainProgress, + *, + deterministic: bool = False, + ) -> dict: + with model.autocast_context: + batch_seed = 0 if deterministic else train_progress.global_step * multi.world_size() + multi.rank() + generator = torch.Generator(device=config.train_device) + generator.manual_seed(batch_seed) + rand = Random(batch_seed) + + latent_image = batch['latent_image'] # (B, 32, H_lat, W_lat) + batch_size = latent_image.shape[0] + latent_height = latent_image.shape[-2] + latent_width = latent_image.shape[-1] + grid_h = latent_height // 2 + grid_w = latent_width // 2 + num_image_tokens = grid_h * grid_w + + text_encoder_output, text_lengths = model.encode_text( + train_device=self.train_device, + batch_size=batch_size, + rand=rand, + tokens=batch.get('tokens'), + tokens_mask=batch.get('tokens_mask'), + text_encoder_output=batch.get('text_encoder_hidden_state'), + text_encoder_dropout_probability=config.text_encoder.dropout_probability if not deterministic else None, + ) + max_text_tokens = text_encoder_output.shape[1] + + # patchify [B, 32, H, W] -> packed (B, num_image_tokens, 128), then bn-normalize in packed space (the + # sampler scales the packed sequence the same way). + packed_latent_image = model.patchify_latents(latent_image.float()) + scaled_latent_image = model.scale_latents(packed_latent_image) + + latent_noise = self._create_noise(scaled_latent_image, config, generator) + + shift = model.calculate_timestep_shift(latent_height, latent_width) + timestep = self._get_timestep_discrete( + model.noise_scheduler.config['num_train_timesteps'], + deterministic, + generator, + batch_size, + config, + shift=shift if config.dynamic_timestep_shifting else config.timestep_shift, + ) + + scaled_noisy_latent_image, sigma = self._add_noise_discrete( + scaled_latent_image, + latent_noise, + timestep, + model.noise_scheduler.timesteps, + ) + + # build the packed [left-pad][text][image] layout (shared helper; identical to the sampler) + position_ids, segment_ids, indicator = model.prepare_packed_ids( + text_lengths, grid_h, grid_w, max_text_tokens, self.train_device, + ) + + # encode_text already returns features left-aligned to this layout, so pack them directly + dtype = model.train_dtype.torch_dtype() + llm_features = model.pack_llm_features(text_encoder_output, num_image_tokens).to(dtype) + + # hidden states: zero latents over the text positions, the noisy image latents over the image positions + text_z_padding = torch.zeros( + batch_size, max_text_tokens, scaled_noisy_latent_image.shape[-1], + dtype=torch.float32, device=self.train_device, + ) + hidden_states = torch.cat([text_z_padding, scaled_noisy_latent_image], dim=1).to(dtype) + + # the transformer's timestep is flow-matching time in [0, 1] (0 = noise, 1 = data), not the discrete index; + # sigma is the noise fraction from _add_noise_discrete, so model time = 1 - sigma (matches the sampler). + model_time = (1.0 - sigma).reshape(batch_size) + + predicted = model.transformer( + hidden_states=hidden_states, + timestep=model_time, + encoder_hidden_states=llm_features, + position_ids=position_ids, + segment_ids=segment_ids, + indicator=indicator, + return_dict=False, + )[0] + predicted_flow = predicted[:, max_text_tokens:].float() + + # The transformer's velocity convention is data - noise. + flow = scaled_latent_image - latent_noise + model_output_data = { + 'loss_type': 'target', + 'timestep': timestep, + # unpatchify both back to (B, 32, H, W) to match the latent mask shape for masked training + 'predicted': model.unpatchify_latents(predicted_flow, grid_h, grid_w), + 'target': model.unpatchify_latents(flow, grid_h, grid_w), + } + + if config.debug_mode: + with torch.no_grad(): + predicted_scaled_latent_image = scaled_noisy_latent_image + predicted_flow * sigma + self._save_tokens('7-prompt', batch['tokens'], model.tokenizer, config, train_progress) + self._save_latent('1-noise', model.unpatchify_latents(latent_noise, grid_h, grid_w), config, train_progress) + self._save_latent('2-noisy_image', model.unpatchify_latents(scaled_noisy_latent_image, grid_h, grid_w), config, train_progress) + self._save_latent('3-predicted_flow', model.unpatchify_latents(predicted_flow, grid_h, grid_w), config, train_progress) + self._save_latent('4-flow', model.unpatchify_latents(flow, grid_h, grid_w), config, train_progress) + self._save_latent('5-predicted_image', model.unpatchify_latents(predicted_scaled_latent_image, grid_h, grid_w), config, train_progress) + self._save_latent('6-image', model.unpatchify_latents(scaled_latent_image, grid_h, grid_w), config, train_progress) + + return model_output_data + + def calculate_loss( + self, + model: IdeogramModel, + batch: dict, + data: dict, + config: TrainConfig, + ) -> Tensor: + return self._flow_matching_losses( + batch=batch, + data=data, + config=config, + train_device=self.train_device, + sigmas=model.noise_scheduler.sigmas, + ).mean() + + def prepare_text_caching(self, model: IdeogramModel, config: TrainConfig): + model.to(self.temp_device) + model.text_encoder_to(self.train_device) + model.eval() + torch_gc() diff --git a/modules/modelSetup/IdeogramFineTuneSetup.py b/modules/modelSetup/IdeogramFineTuneSetup.py new file mode 100644 index 000000000..a259af03e --- /dev/null +++ b/modules/modelSetup/IdeogramFineTuneSetup.py @@ -0,0 +1,96 @@ +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSetup.BaseIdeogramSetup import BaseIdeogramSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.ModuleFilter import ModuleFilter +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +@factory.register(BaseModelSetup, ModelType.IDEOGRAM_4, TrainingMethod.FINE_TUNE) +class IdeogramFineTuneSetup( + BaseIdeogramSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: IdeogramModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + self._create_model_part_parameters( + parameter_group_collection, "transformer", model.transformer, config.transformer, + freeze=ModuleFilter.create(config), debug=config.debug_mode, + ) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: IdeogramModel, + config: TrainConfig, + ): + self._setup_model_part_requires_grad("transformer", model.transformer, config.transformer, model.train_progress) + model.vae.requires_grad_(False) + model.text_encoder.requires_grad_(False) + # the unconditional transformer is never trained (only the negative branch of the dual-network CFG at + # sampling), and may not be loaded at all + if model.unconditional_transformer is not None: + model.unconditional_transformer.requires_grad_(False) + + def setup_model( + self, + model: IdeogramModel, + config: TrainConfig, + ): + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: IdeogramModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + # the unconditional transformer is only needed for sampling; keep it off the train device during training + model.unconditional_transformer_to(self.temp_device) + + model.text_encoder.eval() + model.vae.eval() + if model.unconditional_transformer is not None: + model.unconditional_transformer.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: IdeogramModel, + config: TrainConfig, + train_progress: TrainProgress, + ): + self.__setup_requires_grad(model, config) diff --git a/modules/modelSetup/IdeogramLoRASetup.py b/modules/modelSetup/IdeogramLoRASetup.py new file mode 100644 index 000000000..e3f83d8bf --- /dev/null +++ b/modules/modelSetup/IdeogramLoRASetup.py @@ -0,0 +1,105 @@ +from modules.model.IdeogramModel import IdeogramModel +from modules.modelSetup.BaseIdeogramSetup import BaseIdeogramSetup +from modules.modelSetup.BaseModelSetup import BaseModelSetup +from modules.module.LoRAModule import LoRAModuleWrapper +from modules.util import factory +from modules.util.config.TrainConfig import TrainConfig +from modules.util.enum.ModelType import ModelType +from modules.util.enum.TrainingMethod import TrainingMethod +from modules.util.NamedParameterGroup import NamedParameterGroupCollection +from modules.util.optimizer_util import init_model_parameters +from modules.util.TrainProgress import TrainProgress + +import torch + + +@factory.register(BaseModelSetup, ModelType.IDEOGRAM_4, TrainingMethod.LORA) +class IdeogramLoRASetup( + BaseIdeogramSetup, +): + def __init__( + self, + train_device: torch.device, + temp_device: torch.device, + debug_mode: bool, + ): + super().__init__( + train_device=train_device, + temp_device=temp_device, + debug_mode=debug_mode, + ) + + def create_parameters( + self, + model: IdeogramModel, + config: TrainConfig, + ) -> NamedParameterGroupCollection: + parameter_group_collection = NamedParameterGroupCollection() + self._create_model_part_parameters(parameter_group_collection, "transformer", model.transformer_lora, config.transformer) + return parameter_group_collection + + def __setup_requires_grad( + self, + model: IdeogramModel, + config: TrainConfig, + ): + model.text_encoder.requires_grad_(False) + model.transformer.requires_grad_(False) + if model.unconditional_transformer is not None: + model.unconditional_transformer.requires_grad_(False) + model.vae.requires_grad_(False) + self._setup_model_part_requires_grad("transformer", model.transformer_lora, config.transformer, model.train_progress) + + def setup_model( + self, + model: IdeogramModel, + config: TrainConfig, + ): + model.transformer_lora = LoRAModuleWrapper( + model.transformer, "transformer", config, config.layer_filter.split(","), + fusion_spec=model.fusion_groups(), fuse=config.output_model_format.needs_qkv_fusion(), + ) + + if model.lora_state_dict: + model.transformer_lora.load_state_dict(model.lora_state_dict) + model.lora_state_dict = None + + model.transformer_lora.set_dropout(config.dropout_probability) + model.transformer_lora.to(dtype=config.lora_weight_dtype.torch_dtype()) + model.transformer_lora.hook_to_module() + + params = self.create_parameters(model, config) + self.__setup_requires_grad(model, config) + init_model_parameters(model, params, self.train_device) + + def setup_train_device( + self, + model: IdeogramModel, + config: TrainConfig, + ): + vae_on_train_device = not config.latent_caching + text_encoder_on_train_device = not config.latent_caching + + model.text_encoder_to(self.train_device if text_encoder_on_train_device else self.temp_device) + model.vae_to(self.train_device if vae_on_train_device else self.temp_device) + model.transformer_to(self.train_device) + # the unconditional transformer is only needed for sampling; keep it off the train device during training + model.unconditional_transformer_to(self.temp_device) + + model.text_encoder.eval() + model.vae.eval() + if model.unconditional_transformer is not None: + model.unconditional_transformer.eval() + + if config.transformer.train: + model.transformer.train() + else: + model.transformer.eval() + + def after_optimizer_step( + self, + model: IdeogramModel, + config: TrainConfig, + train_progress: TrainProgress, + ): + self.__setup_requires_grad(model, config) diff --git a/modules/ui/BaseConvertModelUIView.py b/modules/ui/BaseConvertModelUIView.py index b3642a2c4..69cb64925 100644 --- a/modules/ui/BaseConvertModelUIView.py +++ b/modules/ui/BaseConvertModelUIView.py @@ -36,6 +36,7 @@ def build_content(self, frame, controller, ui_state, on_model_or_method_change): ("Anima", ModelType.ANIMA), ("Krea 2", ModelType.KREA_2), ("ZImage", ModelType.Z_IMAGE), + ("Ideogram 4", ModelType.IDEOGRAM_4), ], ui_state, "model_type", command=on_model_or_method_change) # training method diff --git a/modules/ui/BaseModelTabView.py b/modules/ui/BaseModelTabView.py index 0757a71c2..34a9ea8dd 100644 --- a/modules/ui/BaseModelTabView.py +++ b/modules/ui/BaseModelTabView.py @@ -32,6 +32,7 @@ def build_content(self, frame, controller, ui_state): allow_override_prior=model_type.is_stable_cascade(), has_transformer="transformer" in parts, allow_override_transformer=controller.supports_override_transformer(), + has_unconditional_transformer="unconditional_transformer" in parts, has_text_encoder=not model_type.has_multiple_text_encoders(), has_text_encoder_1=model_type.has_multiple_text_encoders(), has_text_encoder_2="text_encoder_2" in parts, @@ -115,6 +116,7 @@ def __create_base_components( allow_override_prior: bool = False, has_transformer: bool = False, allow_override_transformer: bool = False, + has_unconditional_transformer: bool = False, allow_override_text_encoder_4: bool = False, has_text_encoder: bool = False, has_text_encoder_1: bool = False, @@ -168,6 +170,15 @@ def __create_base_components( row += 1 + if has_unconditional_transformer: + # unconditional transformer weight dtype + self.components.label(frame, row, 3, "Unconditional Transformer Data Type", + tooltip="The weight data type of the unconditional transformer, used for the negative branch of CFG during sampling") + self.components.options_kv(frame, row, 4, self.__create_dtype_options(include_a8=True), + ui_state, "unconditional_transformer.weight_dtype") + + row += 1 + presets = controller.get_presets() self.components.label(frame, row, 0, "Quantization") diff --git a/modules/ui/BaseSampleFrameView.py b/modules/ui/BaseSampleFrameView.py index c3e5860c9..eacd8c0ac 100644 --- a/modules/ui/BaseSampleFrameView.py +++ b/modules/ui/BaseSampleFrameView.py @@ -15,8 +15,9 @@ def build_content(self, top_frame, bottom_frame, ui_state, controller, include_p self.components.entry(top_frame, 0, 1, ui_state, "prompt") # negative prompt - self.components.label(top_frame, 1, 0, "negative prompt:") - self.components.entry(top_frame, 1, 1, ui_state, "negative_prompt") + if controller.supports_negative_prompt(): + self.components.label(top_frame, 1, 0, "negative prompt:") + self.components.entry(top_frame, 1, 1, ui_state, "negative_prompt") if include_settings: # width diff --git a/modules/ui/BaseTimestepDistributionWindowView.py b/modules/ui/BaseTimestepDistributionWindowView.py index 82e79ef01..1f29e7079 100644 --- a/modules/ui/BaseTimestepDistributionWindowView.py +++ b/modules/ui/BaseTimestepDistributionWindowView.py @@ -42,5 +42,5 @@ def build_content(self, frame, controller, ui_state): # dynamic timestep shifting self.components.label(frame, 6, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Dynamic Timestep Shifting is not shown in the preview. For Ideogram, the shifting instead follows the model's own resolution-aware sampling schedule. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) self.components.switch(frame, 6, 1, ui_state, "dynamic_timestep_shifting") diff --git a/modules/ui/BaseTrainingTabView.py b/modules/ui/BaseTrainingTabView.py index eef9c82e3..29a54bb1e 100644 --- a/modules/ui/BaseTrainingTabView.py +++ b/modules/ui/BaseTrainingTabView.py @@ -64,6 +64,8 @@ def build(self, column_0, column_1, column_2, controller, ui_state): self.__setup_z_image_ui(column_0, column_1, column_2, controller, ui_state) elif model_type.is_ernie(): self.__setup_ernie_ui(column_0, column_1, column_2, controller, ui_state) + elif model_type.is_ideogram(): + self.__setup_ideogram_ui(column_0, column_1, column_2, controller, ui_state) def __setup_stable_diffusion_ui(self, column_0, column_1, column_2, controller, ui_state): self.__create_base_frame(column_0, 0, controller, ui_state) @@ -232,6 +234,19 @@ def __setup_ernie_ui(self, column_0, column_1, column_2, controller, ui_state): self.__create_loss_frame(column_2, 2, controller, ui_state) self.__create_layer_frame(column_2, 3, controller, ui_state) + def __setup_ideogram_ui(self, column_0, column_1, column_2, controller, ui_state): + self.__create_base_frame(column_0, 0, controller, ui_state) + self.__create_text_encoder_frame(column_0, 1, ui_state, supports_clip_skip=False, supports_training=False, supports_dropout=False) + + self.__create_base2_frame(column_1, 0, controller, ui_state) + self.__create_transformer_frame(column_1, 1, ui_state, supports_guidance_scale=False, supports_force_attention_mask=False) + self.__create_unconditional_transformer_frame(column_1, 2, ui_state) + self.__create_noise_frame(column_1, 3, ui_state, supports_dynamic_timestep_shifting=True) + + self.__create_masked_frame(column_2, 1, ui_state) + self.__create_loss_frame(column_2, 2, controller, ui_state) + self.__create_layer_frame(column_2, 3, controller, ui_state) + def __setup_sana_ui(self, column_0, column_1, column_2, controller, ui_state): self.__create_base_frame(column_0, 0, controller, ui_state) self.__create_text_encoder_frame(column_0, 1, ui_state) @@ -449,7 +464,7 @@ def __create_offloading_widgets(self, frame, row, ui_state, part, supports_check return row def __create_text_encoder_frame(self, master, row, ui_state, supports_clip_skip=True, supports_training=True, - supports_sequence_length=False, supports_layer_offloading=True): + supports_sequence_length=False, supports_dropout=True, supports_layer_offloading=True): frame = self.components.section_frame(master, row) row = 0 @@ -466,11 +481,12 @@ def __create_text_encoder_frame(self, master, row, ui_state, supports_clip_skip= row = self.__create_offloading_widgets(frame, row, ui_state, "text_encoder", supports_checkpointing=supports_training, supports_layer_offloading=supports_layer_offloading) - # dropout - self.components.label(frame, row, 0, "Caption Dropout Probability", - tooltip="The Probability for dropping the text encoder conditioning") - self.components.entry(frame, row, 1, ui_state, "text_encoder.dropout_probability") - row += 1 + if supports_dropout: + # dropout + self.components.label(frame, row, 0, "Caption Dropout Probability", + tooltip="The Probability for dropping the text encoder conditioning") + self.components.entry(frame, row, 1, ui_state, "text_encoder.dropout_probability") + row += 1 if supports_training: # train text encoder epochs @@ -681,6 +697,20 @@ def __create_transformer_frame(self, master, row, ui_state, supports_guidance_sc self.components.entry(frame, row, 1, ui_state, "transformer.guidance_scale") row += 1 + def __create_unconditional_transformer_frame(self, master, row, ui_state): + frame = self.components.section_frame(master, row) + row = 0 + + # include unconditional transformer + self.components.label(frame, row, 0, "Include Unconditional Transformer", + tooltip="Loads the dedicated unconditional transformer used for the negative branch of CFG " + "during sampling. If disabled, CFG above 1.0 still works by running an empty prompt " + "through the conditional transformer instead, at reduced VRAM and load time") + self.components.switch(frame, row, 1, ui_state, "unconditional_transformer.include") + row += 1 + + row = self.__create_offloading_widgets(frame, row, ui_state, "unconditional_transformer", supports_checkpointing=False) + def __create_noise_frame(self, master, row, ui_state, supports_generalized_offset_noise: bool = False, supports_dynamic_timestep_shifting: bool = False): @@ -739,7 +769,7 @@ def __create_noise_frame(self, master, row, ui_state, if supports_dynamic_timestep_shifting: # dynamic timestep shifting self.components.label(frame, 9, 0, "Dynamic Timestep Shifting", - tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) + tooltip="Dynamically shift the timestep distribution based on resolution. If enabled, the shifting parameters are taken from the model's scheduler configuration and Timestep Shift is ignored. For Ideogram, the shifting instead follows the model's own resolution-aware sampling schedule. Note: For Z-Image, the dynamic shifting parameters are likely wrong and unknown. Use with care or set your own, fixed shift.", wide_tooltip=True) self.components.switch(frame, 9, 1, ui_state, "dynamic_timestep_shifting") def __create_masked_frame(self, master, row, ui_state): diff --git a/modules/ui/SampleFrameController.py b/modules/ui/SampleFrameController.py index 474c52ab8..4a6e3498b 100644 --- a/modules/ui/SampleFrameController.py +++ b/modules/ui/SampleFrameController.py @@ -15,3 +15,6 @@ def is_inpainting_model(self) -> bool: def is_video_model(self) -> bool: return self.model_type.is_video_model() + + def supports_negative_prompt(self) -> bool: + return self.model_type.supports_negative_prompt() diff --git a/modules/ui/TopBarController.py b/modules/ui/TopBarController.py index 01d1ee0da..052614a73 100644 --- a/modules/ui/TopBarController.py +++ b/modules/ui/TopBarController.py @@ -43,6 +43,7 @@ def get_model_types(self) -> list[tuple[str, ModelType]]: ("Krea 2", ModelType.KREA_2), ("Z-Image", ModelType.Z_IMAGE), ("Ernie Image", ModelType.ERNIE), + ("Ideogram 4", ModelType.IDEOGRAM_4), ] def get_training_methods(self, model_type: ModelType) -> list[tuple[str, TrainingMethod]]: diff --git a/modules/util/ModelNames.py b/modules/util/ModelNames.py index 8dec9a9bf..a4b7310ff 100644 --- a/modules/util/ModelNames.py +++ b/modules/util/ModelNames.py @@ -25,6 +25,7 @@ def __init__( include_text_encoder_2: bool = True, include_text_encoder_3: bool = True, include_text_encoder_4: bool = True, + include_unconditional_transformer: bool = True, ): self.base_model = base_model self.prior_model = prior_model @@ -40,6 +41,7 @@ def __init__( self.include_text_encoder_2 = include_text_encoder_2 self.include_text_encoder_3 = include_text_encoder_3 self.include_text_encoder_4 = include_text_encoder_4 + self.include_unconditional_transformer = include_unconditional_transformer def all_embedding(self): if self.embedding is not None: diff --git a/modules/util/ModelWeightDtypes.py b/modules/util/ModelWeightDtypes.py index 3893b3a5f..dbfdfb74f 100644 --- a/modules/util/ModelWeightDtypes.py +++ b/modules/util/ModelWeightDtypes.py @@ -11,6 +11,7 @@ def __init__( unet: DataType, prior: DataType, transformer: DataType, + unconditional_transformer: DataType, text_encoder: DataType, text_encoder_2: DataType, text_encoder_3: DataType, @@ -29,6 +30,7 @@ def __init__( self.unet = unet self.prior = prior self.transformer = transformer + self.unconditional_transformer = unconditional_transformer self.text_encoder = text_encoder self.text_encoder_2 = text_encoder_2 self.text_encoder_3 = text_encoder_3 @@ -46,6 +48,7 @@ def all_dtypes(self) -> list: self.unet, self.prior, self.transformer, + self.unconditional_transformer, self.text_encoder, self.text_encoder_2, self.text_encoder_3, diff --git a/modules/util/checkpointing_util.py b/modules/util/checkpointing_util.py index 31b3819ea..1f669e9a5 100644 --- a/modules/util/checkpointing_util.py +++ b/modules/util/checkpointing_util.py @@ -101,7 +101,7 @@ def forward(self, *args, **kwargs): return self.__orig(*args, **kwargs) class OffloadCheckpointLayer(BaseCheckpointLayer): - def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, conductor: LayerOffloadConductor, layer_index: int): + def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.device, conductor: LayerOffloadConductor, layer_index: int, checkpointing: bool): super().__init__() assert (orig_module is None or orig_forward is None) and not (orig_module is None and orig_forward is None) @@ -111,6 +111,7 @@ def __init__(self, orig_module: nn.Module, orig_forward, train_device: torch.dev self.dummy = torch.zeros((1,), device=train_device, requires_grad=True) self.conductor = conductor self.layer_index = layer_index + self.checkpointing = checkpointing def __deepcopy__(self, memo): # conductor holds torch.cuda.Stream/Event objects that cannot be deep-copied or pickled. @@ -145,6 +146,12 @@ def forward(self, *args, **kwargs): call_id = _generate_call_index() args = _kwargs_to_args(self.orig_forward if self.checkpoint is None else self.checkpoint.forward, args, kwargs) if torch.is_grad_enabled(): + # a backward will flow through this layer (grad enabled), so offloading needs use_reentrant=True + # checkpointing to move the offloaded tensors back during recompute. Fail loud rather than silently + # enabling checkpointing the part disabled. Under no_grad (e.g. sampling a frozen part) the branch + # below offloads without checkpointing. + if not self.checkpointing: + raise NotImplementedError("offloading requires gradient checkpointing") return torch.utils.checkpoint.checkpoint( self.__checkpointing_forward, self.dummy, @@ -178,20 +185,19 @@ def create_checkpoint( conductor.add_layer(orig_module, included_offload_param_indices) if conductor is not None and conductor.offload_activated(): - # offloading is structurally coupled to use_reentrant=True checkpointing during the back pass: - # the recompute is the only thing firing before_layer/after_layer in the backward direction, so - # both layer and activation offloading need checkpointing to move tensors back for backward. - # Rather than silently forcing checkpointing on when the part disabled it, reject the combination. - if not checkpointing: - raise NotImplementedError("offloading currently requires gradient checkpointing") + # offloading is structurally coupled to use_reentrant=True checkpointing during the back pass: the + # recompute is the only thing firing before_layer/after_layer in the backward direction, so both layer + # and activation offloading need checkpointing to move tensors back for backward. That coupling only + # matters when a backward actually flows, so OffloadCheckpointLayer.forward enforces it per call (fail + # loud under grad, offload freely under no_grad) instead of rejecting the frozen/inference case here. if compile: - layer = OffloadCheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, conductor=conductor, layer_index=layer_index) + layer = OffloadCheckpointLayer(orig_module=orig_module, orig_forward=None, train_device=train_device, conductor=conductor, layer_index=layer_index, checkpointing=checkpointing) #don't compile the checkpointing layer - offloading cannot be compiled: orig_module.compile(fullgraph=True) return layer else: #only patch forward() if possible. Inserting layers is necessary for torch.compile, but causes issues with at least 1 text encoder model. we don't compile text encoders - layer = OffloadCheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, conductor=conductor, layer_index=layer_index) + layer = OffloadCheckpointLayer(orig_module=None, orig_forward=orig_module.forward, train_device=train_device, conductor=conductor, layer_index=layer_index, checkpointing=checkpointing) orig_module.forward = layer.forward return orig_module else: @@ -247,6 +253,14 @@ def enable_checkpointing( conductor = LayerOffloadConductor(model, config, part) if offload else None checkpointing = part.checkpointing_enabled() + # a trained part always has grad flowing through it, so offloading without checkpointing is guaranteed to hit + # OffloadCheckpointLayer.forward's fail-loud path. Reject it here so the misconfiguration surfaces at setup + # instead of the first training step. Frozen parts (part.train == False) are left to the per-call check: e.g. + # Ideogram's unconditional transformer runs only under no_grad during sampling, so it offloads without + # checkpointing there, while a frozen denoiser/TE still fails loud when a trained embedding routes grad through it. + if offload and not checkpointing and part.train: + raise NotImplementedError("offloading requires gradient checkpointing") + layer_index = 0 for type_or_list, param_names in lists: @@ -460,6 +474,16 @@ def enable_checkpointing_for_ernie_transformer( (model.layers, ["x"]), ]) + +def enable_checkpointing_for_ideogram_transformer( + model: nn.Module, + config: TrainConfig, + part: TrainModelPartConfig, +) -> LayerOffloadConductor | None: + return enable_checkpointing(model, config, part, config.compile, [ + (model.layers, ["hidden_states"]), + ]) + def enable_checkpointing_for_krea2_transformer( model: nn.Module, config: TrainConfig, diff --git a/modules/util/config/SampleConfig.py b/modules/util/config/SampleConfig.py index 613b6acbe..9b2b2c0b1 100644 --- a/modules/util/config/SampleConfig.py +++ b/modules/util/config/SampleConfig.py @@ -143,6 +143,16 @@ def _get_model_defaults(model_type) -> dict: "diffusion_steps": 25, "cfg_scale": 4.0, }) + elif model_type.is_ideogram(): + # Ideogram 4 recommends 48 flow-matching steps on a logit-normal schedule with + # guidance held at 7.0 for the main steps (dropping to 3.0 for the final polish steps). + # Lowered to 25 steps and 3.0 guidance here for faster in-training sampling previews. + defaults.update({ + "width": 1024, + "height": 1024, + "diffusion_steps": 25, + "cfg_scale": 3.0, + }) return defaults diff --git a/modules/util/config/TrainConfig.py b/modules/util/config/TrainConfig.py index da7c936cc..deb632f25 100644 --- a/modules/util/config/TrainConfig.py +++ b/modules/util/config/TrainConfig.py @@ -477,6 +477,7 @@ class TrainConfig(BaseConfig): # transformer transformer: TrainModelPartConfig + unconditional_transformer: TrainModelPartConfig quantization: QuantizationConfig # text encoder @@ -874,6 +875,7 @@ def weight_dtypes(self) -> ModelWeightDtypes: self.unet.weight_dtype, self.prior.weight_dtype, self.transformer.weight_dtype, + self.unconditional_transformer.weight_dtype, self.text_encoder.weight_dtype, self.text_encoder_2.weight_dtype, self.text_encoder_3.weight_dtype, @@ -905,6 +907,7 @@ def model_names(self) -> ModelNames: include_text_encoder_2=self.text_encoder_2.include, include_text_encoder_3=self.text_encoder_3.include, include_text_encoder_4=self.text_encoder_4.include, + include_unconditional_transformer=self.unconditional_transformer.include, ) def train_any_embedding(self) -> bool: @@ -1123,6 +1126,13 @@ def default_values() -> 'TrainConfig': transformer.learning_rate = None data.append(("transformer", transformer, TrainModelPartConfig, False)) + unconditional_transformer = TrainModelPartConfig.default_values() + unconditional_transformer.model_name = "" + unconditional_transformer.train = False + unconditional_transformer.gradient_checkpointing = False + unconditional_transformer.activation_offloading = False + data.append(("unconditional_transformer", unconditional_transformer, TrainModelPartConfig, False)) + #quantization layer filter quantization = QuantizationConfig.default_values() data.append(("quantization", quantization, QuantizationConfig, False)) diff --git a/modules/util/enum/ModelType.py b/modules/util/enum/ModelType.py index 1e927a69f..727d71cfd 100644 --- a/modules/util/enum/ModelType.py +++ b/modules/util/enum/ModelType.py @@ -47,6 +47,8 @@ class ModelType(Enum): ERNIE = 'ERNIE' + IDEOGRAM_4 = 'IDEOGRAM_4' + def __str__(self): return self.value @@ -124,6 +126,14 @@ def is_z_image(self): def is_ernie(self): return self == ModelType.ERNIE + def is_ideogram(self): + return self == ModelType.IDEOGRAM_4 + + def supports_negative_prompt(self) -> bool: + # asymmetric dual-network CFG models drive the negative branch from a frozen unconditional network (or an + # empty prompt), not a user-supplied negative prompt + return not self.is_ideogram() + def has_mask_input(self) -> bool: return self == ModelType.STABLE_DIFFUSION_15_INPAINTING \ or self == ModelType.STABLE_DIFFUSION_20_INPAINTING \ @@ -171,7 +181,8 @@ def is_flow_matching(self) -> bool: or self.is_hunyuan_video() \ or self.is_hi_dream() \ or self.is_z_image() \ - or self.is_ernie() + or self.is_ernie() \ + or self.is_ideogram() def is_video_model(self) -> bool: return self.is_hunyuan_video() #incase we add more video models in the future @@ -193,7 +204,7 @@ def supported_training_methods(self) -> tuple[TrainingMethod, ...]: or self.is_chroma(): return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA, TrainingMethod.EMBEDDING) if self.is_qwen() or self.is_z_image() or self.is_flux_2() or self.is_ernie() \ - or self.is_anima() or self.is_krea2(): + or self.is_anima() or self.is_krea2() or self.is_ideogram(): return (TrainingMethod.FINE_TUNE, TrainingMethod.LORA) raise ValueError(f"No supported training methods defined for model type {self}") @@ -234,7 +245,7 @@ def supported_full_model_formats(self) -> list[ModelFormat]: formats.append(ModelFormat.ORIGINAL_SINGLE_FILE) elif (self.is_flux_1() or self.is_flux_2() or self.is_chroma() or self.is_hunyuan_video() or self.is_hi_dream() or self.is_pixart() or self.is_qwen() or self.is_ernie() - or self.is_z_image() or self.is_anima() or self.is_krea2()): + or self.is_z_image() or self.is_anima() or self.is_krea2() or self.is_ideogram()): formats.append(ModelFormat.ORIGINAL_TRANSFORMER) if self.is_z_image(): formats.append(ModelFormat.COMFY_TRANSFORMER) @@ -299,6 +310,7 @@ def supported_output_formats(self, training_method: TrainingMethod) -> list[Mode ModelType.KREA_2: ("transformer", "text_encoder", "vae"), ModelType.Z_IMAGE: ("transformer", "text_encoder", "vae"), ModelType.ERNIE: ("transformer", "text_encoder", "vae"), + ModelType.IDEOGRAM_4: ("transformer", "text_encoder", "unconditional_transformer", "vae"), } diff --git a/modules/util/optimizer/muon_util.py b/modules/util/optimizer/muon_util.py index f819b4ed6..092070d45 100644 --- a/modules/util/optimizer/muon_util.py +++ b/modules/util/optimizer/muon_util.py @@ -46,7 +46,7 @@ def build_muon_adam_key_fn( 'double_stream_blocks', 'single_stream_blocks', ] - case ModelType.Z_IMAGE | ModelType.ERNIE: + case ModelType.Z_IMAGE | ModelType.ERNIE | ModelType.IDEOGRAM_4: default_patterns = [ 'layers', 'refiner', diff --git a/requirements-global.txt b/requirements-global.txt index da3f88c33..b122a90f1 100644 --- a/requirements-global.txt +++ b/requirements-global.txt @@ -32,7 +32,7 @@ pooch==1.8.2 open-clip-torch==2.32.0 # data loader --e git+https://github.com/Nerogar/mgds.git@bae73e5#egg=mgds +-e git+https://github.com/Nerogar/mgds.git@bec5ef2#egg=mgds # optimizers dadaptation==3.2 # dadaptation optimizers diff --git a/resources/sd_model_spec/ideogram_4-lora.json b/resources/sd_model_spec/ideogram_4-lora.json new file mode 100644 index 000000000..31043e15a --- /dev/null +++ b/resources/sd_model_spec/ideogram_4-lora.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Ideogram/lora", + "modelspec.implementation": "https://github.com/huggingface/diffusers", + "modelspec.title": "Ideogram 4 LoRA" +} diff --git a/resources/sd_model_spec/ideogram_4.json b/resources/sd_model_spec/ideogram_4.json new file mode 100644 index 000000000..f38d8ef5e --- /dev/null +++ b/resources/sd_model_spec/ideogram_4.json @@ -0,0 +1,6 @@ +{ + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "Ideogram", + "modelspec.implementation": "https://github.com/huggingface/diffusers", + "modelspec.title": "Ideogram 4" +} diff --git a/training_presets/Chroma1/#chroma Finetune 16GB.json b/training_presets/Chroma/#chroma Finetune 16GB.json similarity index 100% rename from training_presets/Chroma1/#chroma Finetune 16GB.json rename to training_presets/Chroma/#chroma Finetune 16GB.json diff --git a/training_presets/Chroma1/#chroma Finetune 24GB.json b/training_presets/Chroma/#chroma Finetune 24GB.json similarity index 100% rename from training_presets/Chroma1/#chroma Finetune 24GB.json rename to training_presets/Chroma/#chroma Finetune 24GB.json diff --git a/training_presets/Chroma1/#chroma Finetune 8GB.json b/training_presets/Chroma/#chroma Finetune 8GB.json similarity index 100% rename from training_presets/Chroma1/#chroma Finetune 8GB.json rename to training_presets/Chroma/#chroma Finetune 8GB.json diff --git a/training_presets/Chroma1/#chroma LoRA 16GB.json b/training_presets/Chroma/#chroma LoRA 16GB.json similarity index 100% rename from training_presets/Chroma1/#chroma LoRA 16GB.json rename to training_presets/Chroma/#chroma LoRA 16GB.json diff --git a/training_presets/Chroma1/#chroma LoRA 24GB.json b/training_presets/Chroma/#chroma LoRA 24GB.json similarity index 100% rename from training_presets/Chroma1/#chroma LoRA 24GB.json rename to training_presets/Chroma/#chroma LoRA 24GB.json diff --git a/training_presets/Chroma1/#chroma LoRA 8GB.json b/training_presets/Chroma/#chroma LoRA 8GB.json similarity index 100% rename from training_presets/Chroma1/#chroma LoRA 8GB.json rename to training_presets/Chroma/#chroma LoRA 8GB.json diff --git a/training_presets/Ideogram 4/#ideogram Finetune 16GB.json b/training_presets/Ideogram 4/#ideogram Finetune 16GB.json new file mode 100644 index 000000000..847fb58ed --- /dev/null +++ b/training_presets/Ideogram 4/#ideogram Finetune 16GB.json @@ -0,0 +1,55 @@ +{ + "base_model_name": "CalamitousFelicitousness/Ideogram-4-bf16-Diffusers", + "batch_size": 2, + "learning_rate": 1e-5, + "model_type": "IDEOGRAM_4", + "output_model_format": "ORIGINAL_TRANSFORMER", + "resolution": "512", + "compile": true, + "dataloader_threads": 1, + "transformer": { + "train": true, + "weight_dtype": "BFLOAT_16", + "offload_fraction": 0.6 + }, + "unconditional_transformer": { + "weight_dtype": "NFLOAT_4", + "offload_fraction": 1.0 + }, + "text_encoder": { + "train": false, + "weight_dtype": "NFLOAT_4" + }, + "training_method": "FINE_TUNE", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "layer_filter": "layers", + "layer_filter_preset": "blocks", + "quantization": { + "layer_filter": "layers", + "layer_filter_preset": "blocks" + }, + "timestep_distribution": "LOGIT_NORMAL", + "optimizer": { + "optimizer": "ADAFACTOR" + }, + "optimizer_defaults": { + "ADAFACTOR": { + "optimizer": "ADAFACTOR", + "fused_back_pass": true, + "beta1": null, + "clip_threshold": 1.0, + "decay_rate": -0.8, + "eps": 1e-30, + "eps2": 0.001, + "relative_step": false, + "scale_parameter": false, + "stochastic_rounding": true, + "warmup_init": false, + "weight_decay": 0.0 + } + } +} diff --git a/training_presets/Ideogram 4/#ideogram LoRA 16GB.json b/training_presets/Ideogram 4/#ideogram LoRA 16GB.json new file mode 100644 index 000000000..ebe109175 --- /dev/null +++ b/training_presets/Ideogram 4/#ideogram LoRA 16GB.json @@ -0,0 +1,36 @@ +{ + "base_model_name": "CalamitousFelicitousness/Ideogram-4-bf16-Diffusers", + "batch_size": 2, + "learning_rate": 5e-5, + "model_type": "IDEOGRAM_4", + "output_model_format": "DIFFUSERS_LORA", + "resolution": "512", + "compile": true, + "transformer": { + "train": true, + "weight_dtype": "INT_W8A8", + "offload_fraction": 0.1 + }, + "unconditional_transformer": { + "weight_dtype": "NFLOAT_4", + "offload_fraction": 1.0 + }, + "text_encoder": { + "train": false, + "weight_dtype": "NFLOAT_4" + }, + "training_method": "LORA", + "vae": { + "weight_dtype": "FLOAT_32" + }, + "train_dtype": "BFLOAT_16", + "output_dtype": "BFLOAT_16", + "layer_filter": "attention,feed_forward", + "layer_filter_preset": "attn-mlp", + "quantization": { + "layer_filter": "layers", + "layer_filter_preset": "blocks" + }, + "timestep_distribution": "LOGIT_NORMAL", + "dataloader_threads": 1 +}