def __init__(self):

super().__init__()

self.sample_prompts_te_outputs = None

self.is_swapping_blocks = False

self._use_unsloth_offload_checkpointing = False

# BatchNorm statistics of the FLUX.2 VAE, kept so that the latents can be normalized after the VAE is freed

self.bn_running_mean: Optional[torch.Tensor] = None

self.bn_running_var: Optional[torch.Tensor] = None

# effective ConvRot INT8 state after loading (pre-quantized files are detected independently of the options)

self.dit_convrot_int8 = False

self.text_encoder_convrot_int8 = False

def assert_extra_args(

self,

args,

train_dataset_group: Union[DatasetGroup, MinimalDataset],

val_dataset_group: Optional[DatasetGroup],

):

flux_train_utils.log_timestep_sampling_info(args)

assert args.text_encoder is not None, "--text_encoder (Qwen3.5-9B checkpoint) is required / --text_encoder（Qwen3.5-9B のチェックポイント）を指定してください"

assert args.vae is not None, "--vae (FLUX.2 VAE checkpoint) is required / --vae（FLUX.2 VAE のチェックポイント）を指定してください"

if args.fp8_base or args.fp8_base_unet:

logger.warning(

"fp8_base and fp8_base_unet are not supported for MOSS-Image: use --convrot_int8 (and --convrot_int8_text_encoder)"

" instead. They are disabled."

" / fp8_base と fp8_base_unet は MOSS-Image ではサポートされていません。代わりに --convrot_int8（と"

" --convrot_int8_text_encoder）を使ってください。無効化します。"

)

args.fp8_base = False

args.fp8_base_unet = False

args.fp8_scaled = False

if args.convrot_int8_bwd == "int8":

# the triton requirement is also checked when the layers are patched; fail early here

assert convrot_int8_kernels.HAS_TRITON, (

"--convrot_int8_bwd int8 requires triton (triton-windows on Windows)"

" / --convrot_int8_bwd int8 には triton（Windows では triton-windows）が必要です"

)

if not (args.convrot_int8 or args.convrot_int8_text_encoder):

logger.info(

"--convrot_int8_bwd int8 only takes effect with a ConvRot INT8 base (--convrot_int8, --convrot_int8_text_encoder"

" or a pre-quantized checkpoint)"

)

if args.mixed_precision == "fp16":

logger.warning(

"mixed_precision fp16 is not recommended for MOSS-Image: the text encoder and the DiT are bf16 models. Use bf16."

" / MOSS-Image では mixed_precision fp16 は推奨されません。Text Encoder と DiT は bf16 のモデルです。bf16 を使ってください。"

)

if args.cache_text_encoder_outputs_to_disk and not args.cache_text_encoder_outputs:

logger.warning("cache_text_encoder_outputs_to_disk is enabled, so cache_text_encoder_outputs is also enabled")

args.cache_text_encoder_outputs = True

if args.cache_text_encoder_outputs:

# caption dropout and multiple caption candidates are supported by MossImageTextEncoderOutputsCachingStrategy

assert train_dataset_group.is_text_encoder_output_cacheable(

cache_supports_dropout=True, supports_multiple_caption_candidates=True

), "when caching Text Encoder output, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used"

assert (

args.network_train_unet_only or not args.cache_text_encoder_outputs

), "network for Text Encoder cannot be trained with caching Text Encoder outputs / Text Encoderの出力をキャッシュしながらText Encoderのネットワークを学習することはできません"

if self.is_train_text_encoder(args):

logger.info(

"Training the Text Encoder (Qwen3.5-9B) LoRA. Note: the official recommendation is to keep the Text Encoder frozen."

" / Text Encoder（Qwen3.5-9B）の LoRA を学習します。公式は Text Encoder の凍結を推奨しています。"

)

assert (

args.blocks_to_swap is None or args.blocks_to_swap == 0

) or not args.cpu_offload_checkpointing, "blocks_to_swap is not supported with cpu_offload_checkpointing"

if args.unsloth_offload_checkpointing:

if not args.gradient_checkpointing:

logger.warning("unsloth_offload_checkpointing is enabled, so gradient_checkpointing is also enabled")

args.gradient_checkpointing = True

assert (

not args.cpu_offload_checkpointing

), "Cannot use both --unsloth_offload_checkpointing and --cpu_offload_checkpointing"

assert (

args.blocks_to_swap is None or args.blocks_to_swap == 0

), "blocks_to_swap is not supported with unsloth_offload_checkpointing"

if args.compile:

assert not args.torch_compile, (

"--compile (per-block torch.compile) and --torch_compile (accelerate dynamo) cannot be used together"

" / --compile（ブロック単位torch.compile）と--torch_compile（accelerate dynamo）は併用できません"

)

assert not (args.compile_fullgraph and args.split_attn), (

"--compile_fullgraph cannot be used with --split_attn (split attention uses dynamic control flow)"

" / --compile_fullgraphは--split_attnと併用できません（split attentionは動的な制御フローを使用します）"

)

# FLUX.2 VAE spatial downscale = 8 and patch size = 2

train_dataset_group.verify_bucket_reso_steps(16)

if val_dataset_group is not None:

val_dataset_group.verify_bucket_reso_steps(16)

def load_target_model(self, args, weight_dtype, accelerator):

self.is_swapping_blocks = args.blocks_to_swap is not None and args.blocks_to_swap > 0

# Text encoder (Qwen3.5-9B) is always bf16 (see cast_text_encoder). Loaded to the CPU, moved to the GPU when needed.

logger.info("Loading MOSS-Image text encoder (Qwen3.5-9B)...")

text_encoder = moss_image_utils.load_text_encoder(

args.text_encoder,

dtype=torch.bfloat16,

device="cpu",

disable_mmap=args.disable_mmap_load_safetensors,

convrot_int8=args.convrot_int8_text_encoder,

convrot_int8_bwd=args.convrot_int8_bwd,

calc_device=accelerator.device, # quantize on the GPU, keep the weights on the CPU

)

text_encoder.eval()

self.text_encoder_convrot_int8 = convrot_int8_utils.is_convrot_int8_model(text_encoder)

# VAE (FLUX.2) is always float32 (see cast_vae)

logger.info("Loading FLUX.2 VAE...")

vae = flux2_vae.load_flux2_vae(args.vae, dtype=torch.float32, device="cpu", disable_mmap=args.disable_mmap_load_safetensors)

vae.eval()

self.bn_running_mean = vae.bn.running_mean.detach().clone()

self.bn_running_var = vae.bn.running_var.detach().clone()

# Return format: (model_type, text_encoders, vae, unet)

return "moss_image", [text_encoder], vae, None # DiT is loaded lazily

def load_unet_lazily(self, args, weight_dtype, accelerator, text_encoders) -> tuple[nn.Module, list[nn.Module]]:

loading_device = "cpu" if self.is_swapping_blocks else accelerator.device

attn_mode = "torch"

if args.xformers:

attn_mode = "xformers"

if args.attn_mode is not None:

attn_mode = args.attn_mode

logger.info(f"Loading MOSS-Image DiT with attn_mode={attn_mode}, split_attn: {args.split_attn}...")

model = moss_image_utils.load_moss_image_model(

accelerator.device,

args.pretrained_model_name_or_path,

attn_mode,

args.split_attn,

loading_device,

weight_dtype,

disable_mmap=args.disable_mmap_load_safetensors,

convrot_int8=args.convrot_int8,

convrot_int8_bwd=args.convrot_int8_bwd,

)

self.dit_convrot_int8 = convrot_int8_utils.is_convrot_int8_model(model)

# The base NetworkTrainer only calls enable_gradient_checkpointing(cpu_offload=...); unsloth is re-applied in

# prepare_unet_with_accelerator.

self._use_unsloth_offload_checkpointing = args.unsloth_offload_checkpointing

if self.is_swapping_blocks:

# The base weights are frozen in LoRA training. Freeze them here, before enabling block swap, for the H2D-only offloader.

model.requires_grad_(False)

swap_config = BlockSwapConfig.from_args(args, accelerator.device, supports_backward=True)

logger.info(

f"enable block swap: blocks_to_swap={args.blocks_to_swap}, H2D-only: {swap_config.h2d_only}, "

f"use pinned memory: {swap_config.use_pinned_memory}"

)

model.enable_block_swap(args.blocks_to_swap, swap_config)

return model, text_encoders

# region strategies

def get_tokenize_strategy(self, args):

return strategy_moss_image.MossImageTokenizeStrategy(args.moss_max_token_length, args.tokenizer_cache_dir)

def get_tokenizers(self, tokenize_strategy: strategy_moss_image.MossImageTokenizeStrategy):

return [tokenize_strategy.tokenizer]

def get_latents_caching_strategy(self, args):

return strategy_moss_image.MossImageLatentsCachingStrategy(

args.cache_latents_to_disk, args.vae_batch_size, args.skip_cache_check

)

def get_text_encoding_strategy(self, args):

return strategy_moss_image.MossImageTextEncodingStrategy()

def get_text_encoder_outputs_caching_strategy(self, args):

if args.cache_text_encoder_outputs:

return strategy_moss_image.MossImageTextEncoderOutputsCachingStrategy(

args.cache_text_encoder_outputs_to_disk, args.text_encoder_batch_size, args.skip_cache_check, False

)

return None

# endregion

def post_process_network(self, args, accelerator, network, text_encoders, unet):

pass

def get_models_for_text_encoding(self, args, accelerator, text_encoders):

if args.cache_text_encoder_outputs:

return None # no text encoders needed for encoding

return text_encoders

def cast_text_encoder(self, args):

return False # the text encoder stays in bf16 regardless of the mixed precision setting

def cast_vae(self, args):

return False # the VAE stays in float32

def cast_unet(self, args):

# The DiT is loaded in weight_dtype by load_moss_image_model. `Module.to(dtype)` would also cast the fp32

# `scale_weight` buffers of the ConvRot INT8 layers, so the cast is skipped here.

return False

def cache_text_encoder_outputs_if_needed(

self, args, accelerator: Accelerator, unet, vae, text_encoders, dataset: DatasetGroup, weight_dtype

):

if args.cache_text_encoder_outputs:

if not args.lowram:

# the DiT is not loaded yet, so only the VAE is moved

logger.info("move vae to cpu to save memory")

org_vae_device = vae.device

vae.to("cpu")

clean_memory_on_device(accelerator.device)

logger.info("move text encoder to gpu")

text_encoders[0].to(accelerator.device)

# no autocast: the text encoder is bf16 and the official pipeline runs it in bf16

dataset.new_cache_text_encoder_outputs(text_encoders, accelerator)

# cache sample prompts

if args.sample_prompts is not None:

logger.info(f"cache Text Encoder outputs for sample prompts: {args.sample_prompts}")

tokenize_strategy = strategy_base.TokenizeStrategy.get_strategy()

text_encoding_strategy = strategy_base.TextEncodingStrategy.get_strategy()

prompts = sampling.load_prompts(args.sample_prompts)

sample_prompts_te_outputs = {}

with torch.no_grad():

for prompt_dict in prompts:

for p in [prompt_dict.get("prompt", ""), prompt_dict.get("negative_prompt", "")]:

if p not in sample_prompts_te_outputs:

logger.info(f" cache TE outputs for: {p}")

tokens = tokenize_strategy.tokenize(p)

sample_prompts_te_outputs[p] = text_encoding_strategy.encode_tokens(

tokenize_strategy, text_encoders, tokens

)

self.sample_prompts_te_outputs = sample_prompts_te_outputs

accelerator.wait_for_everyone()

logger.info("move text encoder back to cpu")

text_encoders[0].to("cpu")

if not args.lowram:

logger.info("move vae back to original device")

vae.to(org_vae_device)

clean_memory_on_device(accelerator.device)

else:

# move text encoder to device for encoding during training/validation

text_encoders[0].to(accelerator.device)

def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):

text_encoders = text_encoder if isinstance(text_encoder, list) else [text_encoder] # compatibility

te = self.get_models_for_text_encoding(args, accelerator, text_encoders)

te = te[0] if te is not None else None

text_encoding_strategy = strategy_base.TextEncodingStrategy.get_strategy()

tokenize_strategy = strategy_base.TokenizeStrategy.get_strategy()

moss_image_train_utils.sample_images(

accelerator,

args,

epoch,

global_step,

unet,

vae,

te,

tokenize_strategy,

text_encoding_strategy,

self.sample_prompts_te_outputs,

)

def get_noise_scheduler(self, args: argparse.Namespace, device: torch.device) -> Any:

return sd3_train_utils.FlowMatchEulerDiscreteScheduler(

num_train_timesteps=moss_image_train_utils.NUM_TRAIN_TIMESTEPS, shift=args.discrete_flow_shift

)

def encode_images_to_latents(self, args, vae: flux2_vae.AutoEncoder, images):

# raw latents [B, 32, H/8, W/8] (same as the cache). The VAE is float32: run outside autocast.

with torch.autocast(device_type=vae.device.type, enabled=False):

return vae.encode(images.to(vae.dtype))

def shift_scale_latents(self, args, latents):

# raw latents [B, 32, H/8, W/8] -> DiT input [B, 128, H/16, W/16] (2x2 patchify + BatchNorm normalization)

latents = flux2_vae.patchify_latents(latents)

return flux2_vae.normalize_latents(latents, self.bn_running_mean, self.bn_running_var)

def get_noise_pred_and_target(

self,

args,

accelerator,

noise_scheduler,

latents,

batch,

text_encoder_conds,

unet,

network,

weight_dtype,

train_unet,

is_train=True,

):

dit: moss_image_models.MossImage = unet

noise = torch.randn_like(latents)

# Per-sample timestep sampling offset from custom_attributes: timestep_sampling = { offset = ... }

# Only applied during training; validation uses unbiased sampling for comparable loss.

tso = None

if is_train and "custom_attributes" in batch:

offsets = [ca.get("timestep_sampling", {}).get("offset", 0.0) for ca in batch["custom_attributes"]]

t = torch.tensor(offsets, dtype=torch.float32)

if t.abs().sum() > 0:

tso = t

# Timesteps and sigmas are kept in float32 (bf16 would quantize sigma * 1000 to multiples of 4). The DiT casts

# the noisy input to its own dtype.

noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timesteps(

args, noise_scheduler, latents, noise, accelerator.device, torch.float32, timestep_sampling_offset=tso

)

# timesteps are sigma * 1000: passed to the DiT as is (the official pipeline does not divide by 1000)

prompt_embeds, attn_mask = text_encoder_conds[:2] # caption_dropout_rate (cached outputs) is not needed here

prompt_embeds = prompt_embeds.to(accelerator.device, dtype=weight_dtype)

attn_mask = attn_mask.to(accelerator.device)

# Gradient checkpointing (non-reentrant) needs an input which requires grad

if args.gradient_checkpointing:

noisy_model_input.requires_grad_(True)

prompt_embeds.requires_grad_(True)

with torch.set_grad_enabled(is_train), accelerator.autocast():

model_pred = dit(noisy_model_input, timesteps, prompt_embeds, attn_mask)

# Rectified flow target: v = noise - x0

target = noise - latents

# 1 for logit_normal (the official training uses no loss weighting)

weighting = flux_train_utils.compute_loss_weighting_for_sd3(args.weighting_scheme, sigmas)

return model_pred, target, timesteps, weighting

def process_batch(

self,

batch,

text_encoders,

unet,

network,

vae,

noise_scheduler,

vae_dtype,

weight_dtype,

accelerator,

args,

text_encoding_strategy,

tokenize_strategy,

is_train=True,

train_text_encoder=True,

train_unet=True,

) -> torch.Tensor:

"""Caption dropout for cached text encoder outputs: dropped samples use the cached unconditional embeddings."""

text_encoder_outputs_list = batch.get("text_encoder_outputs_list", None)

if text_encoder_outputs_list is not None:

moss_text_encoding_strategy: strategy_moss_image.MossImageTextEncodingStrategy = text_encoding_strategy

caption_dropout_rates = text_encoder_outputs_list[-1]

# [prompt_embeds, attn_mask, uncond_embeds, uncond_mask] -> [prompt_embeds, attn_mask]

text_encoder_outputs_list = moss_text_encoding_strategy.drop_cached_text_encoder_outputs(

*text_encoder_outputs_list[:-1], caption_dropout_rates=caption_dropout_rates

)

# Add the caption dropout rates back to the list for validation dataset (which is re-used batch items)

batch["text_encoder_outputs_list"] = text_encoder_outputs_list + [caption_dropout_rates]

return super().process_batch(

batch,

text_encoders,

unet,

network,

vae,

noise_scheduler,

vae_dtype,

weight_dtype,

accelerator,

args,

text_encoding_strategy,

tokenize_strategy,

is_train,

train_text_encoder,

train_unet,

)

def post_process_loss(self, loss, args, timesteps, noise_scheduler):

return loss

def get_sai_model_spec(self, args):

return model_io.get_sai_model_spec_dataclass(None, args, False, True, False, moss_image="base").to_metadata_dict()

def update_metadata(self, metadata, args):

metadata["ss_weighting_scheme"] = args.weighting_scheme

metadata["ss_logit_mean"] = args.logit_mean

metadata["ss_logit_std"] = args.logit_std

metadata["ss_mode_scale"] = args.mode_scale

metadata["ss_timestep_sampling"] = args.timestep_sampling

metadata["ss_sigmoid_scale"] = args.sigmoid_scale

metadata["ss_discrete_flow_shift"] = args.discrete_flow_shift

metadata["ss_convrot_int8"] = self.dit_convrot_int8

metadata["ss_convrot_int8_text_encoder"] = self.text_encoder_convrot_int8

def is_text_encoder_not_needed_for_training(self, args):

return args.cache_text_encoder_outputs and not self.is_train_text_encoder(args)

def prepare_text_encoder_grad_ckpt_workaround(self, index, text_encoder):

# Set the first floating point parameter's requires_grad to True to workaround Accelerate gradient checkpointing bug.

# The first parameter (token embedding) is int8 with ConvRot INT8 and cannot require grad.

first_param = next(p for p in text_encoder.parameters() if p.is_floating_point())

first_param.requires_grad_(True)

def prepare_unet_with_accelerator(

self, args: argparse.Namespace, accelerator: Accelerator, unet: torch.nn.Module

) -> torch.nn.Module:

# The base NetworkTrainer only calls enable_gradient_checkpointing(cpu_offload=True/False),

# so we re-apply with unsloth_offload if needed (after base has already enabled it).

if self._use_unsloth_offload_checkpointing and args.gradient_checkpointing:

unet.enable_gradient_checkpointing(unsloth_offload=True)

if not self.is_swapping_blocks:

model = super().prepare_unet_with_accelerator(args, accelerator, unet)

else:

model = accelerator.prepare(unet, device_placement=[not self.is_swapping_blocks])

accelerator.unwrap_model(model).move_to_device_except_swap_blocks(accelerator.device)

accelerator.unwrap_model(model).prepare_block_swap_before_forward()

# CUDA perf switches are independent of torch.compile; apply whenever requested.

compile_utils.apply_cuda_optimizations(args)

if args.compile:

# Per-block torch.compile of the 42 main blocks. Reach the real model via unwrap_model so the underlying

# ModuleList is mutated regardless of any DDP wrapper.

# ConvRot INT8 Linears compile as well: the fused Triton path is a torch.library custom op (opaque node, no

# graph break), so only block swapping needs the Linear exclusion.

dit = accelerator.unwrap_model(model)

compile_utils.compile_transformer(args, dit, [dit.layers], disable_linear=self.is_swapping_blocks)

return model

def on_validation_step_end(self, args, accelerator, network, text_encoders, unet, batch, weight_dtype):

if self.is_swapping_blocks:

# prepare for next forward: because backward pass is not called, we need to prepare it here

parser = train_network.setup_parser()

args_util.add_dit_training_arguments(parser)

moss_image_train_utils.add_moss_image_model_arguments(parser)

moss_image_train_utils.add_moss_image_training_arguments(parser)