onnx / moss_image_train_network.py

Upload folder using huggingface_hub

3251839 verified 24 days ago

24.5 kB

# MOSS-Image LoRA training script

import argparse
from typing import Any, Optional, Union

import torch
import torch.nn as nn
from accelerate import Accelerator
from library.device_utils import init_ipex, clean_memory_on_device
from library.custom_offloading_utils import BlockSwapConfig

init_ipex()

from library import (
    convrot_int8_kernels,
    convrot_int8_utils,
    flux2_vae,
    flux_train_utils,
    moss_image_models,
    moss_image_train_utils,
    moss_image_utils,
    sd3_train_utils,
    strategy_base,
    strategy_moss_image,
    sampling,
)
import library.args as args_util
import library.compile_utils as compile_utils
import library.model_io as model_io
from library.dataset import DatasetGroup, MinimalDataset
import train_network
from library.utils import setup_logging

setup_logging()
import logging

logger = logging.getLogger(__name__)


class MossImageNetworkTrainer(train_network.NetworkTrainer):
    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 are also excluded: the custom autograd.Function + autotuned Triton kernels do not compile.
            dit = accelerator.unwrap_model(model)
            compile_utils.compile_transformer(
                args, dit, [dit.layers], disable_linear=self.is_swapping_blocks or convrot_int8_utils.is_convrot_int8_model(dit)
            )

        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
            accelerator.unwrap_model(unet).prepare_block_swap_before_forward()


def setup_parser() -> argparse.ArgumentParser:
    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)
    return parser


if __name__ == "__main__":
    parser = setup_parser()

    args = parser.parse_args()
    args_util.verify_command_line_training_args(args)
    args = args_util.read_config_from_file(args, parser)

    if args.attn_mode == "sdpa":
        args.attn_mode = "torch"  # backward compatibility

    if args.show_timesteps:
        moss_image_train_utils.show_timesteps(args)
    else:
        trainer = MossImageNetworkTrainer()
        trainer.train(args)

	# MOSS-Image LoRA training script

	import argparse
	from typing import Any, Optional, Union

	import torch
	import torch.nn as nn
	from accelerate import Accelerator
	from library.device_utils import init_ipex, clean_memory_on_device
	from library.custom_offloading_utils import BlockSwapConfig

	init_ipex()

	from library import (
	convrot_int8_kernels,
	convrot_int8_utils,
	flux2_vae,
	flux_train_utils,
	moss_image_models,
	moss_image_train_utils,
	moss_image_utils,
	sd3_train_utils,
	strategy_base,
	strategy_moss_image,
	sampling,
	)
	import library.args as args_util
	import library.compile_utils as compile_utils
	import library.model_io as model_io
	from library.dataset import DatasetGroup, MinimalDataset
	import train_network
	from library.utils import setup_logging

	setup_logging()
	import logging

	logger = logging.getLogger(__name__)


	class MossImageNetworkTrainer(train_network.NetworkTrainer):
	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 are also excluded: the custom autograd.Function + autotuned Triton kernels do not compile.
	dit = accelerator.unwrap_model(model)
	compile_utils.compile_transformer(
	args, dit, [dit.layers], disable_linear=self.is_swapping_blocks or convrot_int8_utils.is_convrot_int8_model(dit)
	)

	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
	accelerator.unwrap_model(unet).prepare_block_swap_before_forward()


	def setup_parser() -> argparse.ArgumentParser:
	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)
	return parser


	if __name__ == "__main__":
	parser = setup_parser()

	args = parser.parse_args()
	args_util.verify_command_line_training_args(args)
	args = args_util.read_config_from_file(args, parser)

	if args.attn_mode == "sdpa":
	args.attn_mode = "torch" # backward compatibility

	if args.show_timesteps:
	moss_image_train_utils.show_timesteps(args)
	else:
	trainer = MossImageNetworkTrainer()
	trainer.train(args)