anima_train_network.py

# Anima LoRA training script

import argparse
import ast
import re
import types
from collections import Counter
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

init_ipex()

from library import (
    anima_models,
    anima_train_utils,
    anima_utils,
    flux_train_utils,
    qwen_image_autoencoder_kl,
    sd3_train_utils,
    strategy_anima,
    strategy_base,
    train_util,
)
import train_network
from library.utils import setup_logging

setup_logging()
import logging

logger = logging.getLogger(__name__)


class AnimaNetworkTrainer(train_network.NetworkTrainer):
    def __init__(self):
        super().__init__()
        self.sample_prompts_te_outputs = None
        self._random_noise_shift_current: Optional[float] = None
        self._random_noise_multiplier_current: Optional[float] = None

    def assert_extra_args(
        self,
        args,
        train_dataset_group: Union[train_util.DatasetGroup, train_util.MinimalDataset],
        val_dataset_group: Optional[train_util.DatasetGroup],
    ):
        if args.random_noise_shift < 0.0:
            raise ValueError("random_noise_shift must be greater than or equal to 0.0")
        if args.random_noise_multiplier < 0.0:
            raise ValueError("random_noise_multiplier must be greater than or equal to 0.0")
        if args.random_noise_shift_decay < 0.0 or args.random_noise_shift_decay > 1.0:
            raise ValueError("random_noise_shift_decay must be between 0.0 and 1.0")
        if args.random_noise_multiplier_decay < 0.0 or args.random_noise_multiplier_decay > 1.0:
            raise ValueError("random_noise_multiplier_decay must be between 0.0 and 1.0")
        if args.knn_noise_k < 0:
            raise ValueError("knn_noise_k must be greater than or equal to 0")

        if args.fp8_base or args.fp8_base_unet:
            logger.warning("fp8_base and fp8_base_unet are not supported. / fp8_baseとfp8_base_unetはサポートされていません。")
            args.fp8_base = False
            args.fp8_base_unet = False
        args.fp8_scaled = False  # Anima DiT does not support fp8_scaled

        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:
            assert train_dataset_group.is_text_encoder_output_cacheable(
                cache_supports_dropout=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のネットワークを学習することはできません"

        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"

        # Bridge Anima's llm_adapter_lr to a network arg, keeping explicit user setting as priority.
        if args.network_module == "lycoris.kohya":
            if args.network_args is None:
                args.network_args = []
            has_train_llm_adapter = any(
                net_arg.split("=", 1)[0].strip() == "train_llm_adapter" for net_arg in args.network_args
            )
            if not has_train_llm_adapter and args.llm_adapter_lr is not None:
                train_llm_adapter = args.llm_adapter_lr > 0
                args.network_args.append(f"train_llm_adapter={'true' if train_llm_adapter else 'false'}")
                logger.info(
                    f"auto-set network_args train_llm_adapter={train_llm_adapter} from llm_adapter_lr={args.llm_adapter_lr}"
                )

        train_dataset_group.verify_bucket_reso_steps(16)  # WanVAE spatial downscale = 8 and patch size = 2
        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

        # Load Qwen3 text encoder (tokenizers already loaded in get_tokenize_strategy)
        logger.info("Loading Qwen3 text encoder...")
        qwen3_text_encoder, _ = anima_utils.load_qwen3_text_encoder(args.qwen3, dtype=weight_dtype, device="cpu")
        qwen3_text_encoder.eval()

        # Load VAE
        logger.info("Loading Anima VAE...")
        vae = qwen_image_autoencoder_kl.load_vae(
            args.vae, device="cpu", disable_mmap=True, spatial_chunk_size=args.vae_chunk_size, disable_cache=args.vae_disable_cache
        )
        vae.to(weight_dtype)
        vae.eval()

        # Return format: (model_type, text_encoders, vae, unet)
        return "anima", [qwen3_text_encoder], vae, None  # unet loaded lazily

    def load_unet_lazily(self, args, weight_dtype, accelerator, text_encoders) -> tuple[nn.Module, list[nn.Module]]:
        loading_dtype = None if args.fp8_scaled else weight_dtype
        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

        # Load DiT
        logger.info(f"Loading Anima DiT model with attn_mode={attn_mode}, split_attn: {args.split_attn}...")
        model = anima_utils.load_anima_model(
            accelerator.device,
            args.pretrained_model_name_or_path,
            attn_mode,
            args.split_attn,
            loading_device,
            loading_dtype,
            args.fp8_scaled,
        )

        # Store unsloth preference so that when the base NetworkTrainer calls
        # dit.enable_gradient_checkpointing(cpu_offload=...), we can override to use unsloth.
        # The base trainer only passes cpu_offload, so we store the flag on the model.
        self._use_unsloth_offload_checkpointing = args.unsloth_offload_checkpointing

        # Block swap
        self.is_swapping_blocks = args.blocks_to_swap is not None and args.blocks_to_swap > 0
        if self.is_swapping_blocks:
            logger.info(f"enable block swap: blocks_to_swap={args.blocks_to_swap}")
            model.enable_block_swap(args.blocks_to_swap, accelerator.device)

        return model, text_encoders

    def get_tokenize_strategy(self, args):
        # Load tokenizers from paths (called before load_target_model, so self.qwen3_tokenizer isn't set yet)
        tokenize_strategy = strategy_anima.AnimaTokenizeStrategy(
            qwen3_path=args.qwen3,
            t5_tokenizer_path=args.t5_tokenizer_path,
            qwen3_max_length=args.qwen3_max_token_length,
            t5_max_length=args.t5_max_token_length,
        )
        return tokenize_strategy

    def get_tokenizers(self, tokenize_strategy: strategy_anima.AnimaTokenizeStrategy):
        return [tokenize_strategy.qwen3_tokenizer]

    def get_latents_caching_strategy(self, args):
        return strategy_anima.AnimaLatentsCachingStrategy(args.cache_latents_to_disk, args.vae_batch_size, args.skip_cache_check)

    def get_text_encoding_strategy(self, args):
        return strategy_anima.AnimaTextEncodingStrategy()

    def post_process_network(self, args, accelerator, network, text_encoders, unet):
        if args.network_module != "lycoris.kohya":
            return
        if not hasattr(network, "unet_loras") or not hasattr(network, "text_encoder_loras"):
            logger.warning("LyCORIS network object has no expected lora lists, skip post-processing.")
            return

        net_kwargs = {}
        if args.network_args is not None:
            for net_arg in args.network_args:
                if "=" not in net_arg:
                    continue
                key, value = net_arg.split("=", 1)
                net_kwargs[key.strip()] = value.strip()

        def parse_bool(v: Optional[str], default: bool = False) -> bool:
            if v is None:
                return default
            return str(v).strip().lower() in ("1", "true", "yes", "on")

        def parse_pattern_list(v: Optional[str], arg_name: str) -> list[str]:
            if v is None:
                return []
            try:
                obj = ast.literal_eval(v)
            except Exception:
                obj = v
            if isinstance(obj, str):
                return [obj]
            if isinstance(obj, (list, tuple)):
                return [str(x) for x in obj]
            logger.warning(f"{arg_name} should be string/list; got {type(obj).__name__}, ignored.")
            return []

        def compile_patterns(patterns: list[str], label: str) -> list[re.Pattern]:
            out = []
            for p in patterns:
                try:
                    out.append(re.compile(p))
                except re.error as e:
                    logger.warning(f"Invalid regex in {label}: {p} ({e})")
            return out

        # Match lora_anima behavior: default exclusion + user exclusion, with include overriding exclusion.
        exclude_patterns = [r".*(_modulation|_norm|_embedder|final_layer).*"]
        exclude_patterns.extend(parse_pattern_list(net_kwargs.get("exclude_patterns"), "exclude_patterns"))
        include_patterns = parse_pattern_list(net_kwargs.get("include_patterns"), "include_patterns")
        verbose = parse_bool(net_kwargs.get("verbose"), False)
        try:
            report_limit = int(net_kwargs.get("report_limit", "0"))
        except ValueError:
            report_limit = 0

        train_llm_adapter = parse_bool(net_kwargs.get("train_llm_adapter"), args.llm_adapter_lr is not None and args.llm_adapter_lr > 0)
        if not train_llm_adapter:
            exclude_patterns.append(r".*llm_adapter.*")

        exclude_re = compile_patterns(exclude_patterns, "exclude_patterns")
        include_re = compile_patterns(include_patterns, "include_patterns")

        unet_name_map = {id(m): n for n, m in unet.named_modules()}
        te_name_map = {}
        if text_encoders is not None:
            tes = text_encoders if isinstance(text_encoders, list) else [text_encoders]
            for te in tes:
                if te is None:
                    continue
                for n, m in te.named_modules():
                    te_name_map[id(m)] = n

        logger.info("LyCORIS post-process report")
        logger.info("  target network module: %s", args.network_module)
        logger.info("  train_llm_adapter: %s", train_llm_adapter)
        logger.info("  default exclude pattern: %s", r".*(_modulation|_norm|_embedder|final_layer).*")
        logger.info("  user exclude patterns: %s", parse_pattern_list(net_kwargs.get("exclude_patterns"), "exclude_patterns"))
        logger.info("  user include patterns: %s", include_patterns)
        logger.info("  network_reg_lrs: %s", net_kwargs.get("network_reg_lrs", "(not set)"))
        if args.network_train_unet_only:
            logger.info("  note: network_train_unet_only=true, effective TE module count is reported as 0")

        def resolve_original_name(lora_obj, is_unet: bool) -> str:
            name_map = unet_name_map if is_unet else te_name_map
            try:
                module_obj = lora_obj.org_module[0]
            except Exception:
                module_obj = None
            if module_obj is not None:
                mapped = name_map.get(id(module_obj))
                if mapped is not None:
                    return mapped
            return getattr(lora_obj, "lora_name", "")

        def should_drop(name: str) -> tuple[bool, str]:
            matched_excludes = [p.pattern for p in exclude_re if p.fullmatch(name)]
            if not matched_excludes:
                return False, ""
            matched_includes = [p.pattern for p in include_re if p.fullmatch(name)]
            if matched_includes:
                return False, ""
            return True, "excluded_by=" + " | ".join(matched_excludes)

        def filter_loras(loras: list, is_unet: bool) -> tuple[list, int, list[tuple[str, str]]]:
            kept = []
            removed = 0
            removed_items: list[tuple[str, str]] = []
            for lora in loras:
                original_name = resolve_original_name(lora, is_unet)
                setattr(lora, "original_name", original_name)
                drop, reason = should_drop(original_name)
                if drop:
                    removed += 1
                    removed_items.append((original_name, reason))
                    if verbose:
                        logger.info("LyCORIS post-filter drop: %s (%s)", original_name, reason)
                    continue
                kept.append(lora)
            return kept, removed, removed_items

        te_before = len(network.text_encoder_loras)
        unet_before = len(network.unet_loras)
        network.text_encoder_loras, te_removed, te_removed_items = filter_loras(network.text_encoder_loras, is_unet=False)
        network.unet_loras, unet_removed, unet_removed_items = filter_loras(network.unet_loras, is_unet=True)
        network.loras = network.text_encoder_loras + network.unet_loras

        if args.network_train_unet_only:
            te_before_log = 0
            te_after_log = 0
            te_removed_log = 0
        else:
            te_before_log = te_before
            te_after_log = len(network.text_encoder_loras)
            te_removed_log = te_removed

        logger.info(
            "LyCORIS post-filter: TE %d -> %d (removed %d), U-Net %d -> %d (removed %d)",
            te_before_log,
            te_after_log,
            te_removed_log,
            unet_before,
            len(network.unet_loras),
            unet_removed,
        )
        total_removed_items = [("Text Encoder", n, r) for n, r in te_removed_items] + [("U-Net", n, r) for n, r in unet_removed_items]
        if total_removed_items:
            logger.info("LyCORIS post-filter changed module count: %d", len(total_removed_items))
            reason_counter = Counter([reason for _, _, reason in total_removed_items])
            for reason, cnt in reason_counter.items():
                logger.info("  changed by reason: %s -> %d modules", reason, cnt)
            if verbose and report_limit > 0:
                preview = total_removed_items[: max(report_limit, 0)]
                for scope, name, reason in preview:
                    logger.info("  changed [%s]: %s (%s)", scope, name, reason)
                if len(total_removed_items) > len(preview):
                    logger.info("  ... and %d more changed modules", len(total_removed_items) - len(preview))
        else:
            logger.info("LyCORIS post-filter changed module count: 0")

        # Optional regex-based LR override by module original name.
        reg_lr_spec = net_kwargs.get("network_reg_lrs")
        if not reg_lr_spec:
            return

        reg_lrs = []
        for pair in reg_lr_spec.split(","):
            pair = pair.strip()
            if not pair:
                continue
            if "=" not in pair:
                logger.warning(f"Invalid network_reg_lrs item: {pair}, expected pattern=lr")
                continue
            pattern, lr_str = pair.split("=", 1)
            pattern = pattern.strip()
            lr_str = lr_str.strip()
            try:
                reg = re.compile(pattern)
                lr_val = float(lr_str)
            except Exception as e:
                logger.warning(f"Invalid network_reg_lrs item: {pair} ({e})")
                continue
            reg_lrs.append((reg, pattern, lr_val))
            logger.info("  network_reg_lrs rule: %s -> lr=%s", pattern, lr_val)

        if not reg_lrs:
            return

        def rule_for_module_name(module_name: str):
            for reg, pattern, lr_val in reg_lrs:
                if reg.fullmatch(module_name):
                    return pattern, lr_val
            return None, None

        all_loras_for_stats = list(network.text_encoder_loras) + list(network.unet_loras)
        rule_hits = Counter()
        matched_module_count = 0
        for lora in all_loras_for_stats:
            module_name = getattr(lora, "original_name", getattr(lora, "lora_name", ""))
            pattern, lr_val = rule_for_module_name(module_name)
            if pattern is not None:
                matched_module_count += 1
                rule_hits[f"{pattern} => {lr_val}"] += 1

        logger.info(
            "network_reg_lrs match summary: matched_modules=%d / total_modules=%d",
            matched_module_count,
            len(all_loras_for_stats),
        )
        if rule_hits:
            for k, v in rule_hits.items():
                logger.info("  network_reg_lrs hit: %s (modules=%d)", k, v)
        else:
            logger.info("  network_reg_lrs hit: none")

        def match_reg_lr(module_name: str, default_lr: Optional[float]) -> Optional[float]:
            _, lr_val = rule_for_module_name(module_name)
            if lr_val is not None:
                return lr_val
            return default_lr

        def build_groups(loras: list, base_lr: Optional[float], plus_ratio: Optional[float], scope: str):
            grouped = {}
            descriptions = []
            for lora in loras:
                module_name = getattr(lora, "original_name", getattr(lora, "lora_name", ""))
                module_lr = match_reg_lr(module_name, base_lr)
                for p_name, param in lora.named_parameters():
                    lr_val = module_lr
                    desc = scope
                    if plus_ratio is not None and "lora_up" in p_name:
                        lr_val = None if lr_val is None else lr_val * plus_ratio
                        desc += " plus"
                    if lr_val is None or lr_val == 0:
                        continue
                    key = (lr_val, desc)
                    if key not in grouped:
                        grouped[key] = []
                    grouped[key].append(param)
            param_groups = []
            for (lr_val, desc), params in grouped.items():
                param_groups.append({"params": params, "lr": lr_val})
                descriptions.append(desc)
            return param_groups, descriptions

        def patched_prepare_optimizer_params(self, text_encoder_lr=None, unet_lr: float = 1e-4, learning_rate=None):
            self.requires_grad_(True)
            all_params = []
            lr_descriptions = []

            te_base_lr = text_encoder_lr if text_encoder_lr is not None else learning_rate
            unet_base_lr = unet_lr if unet_lr is not None else learning_rate

            if self.text_encoder_loras:
                te_ratio = self.loraplus_text_encoder_lr_ratio or self.loraplus_lr_ratio
                params, descriptions = build_groups(self.text_encoder_loras, te_base_lr, te_ratio, "textencoder")
                all_params.extend(params)
                lr_descriptions.extend(descriptions)

            if self.unet_loras:
                unet_ratio = self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio
                params, descriptions = build_groups(self.unet_loras, unet_base_lr, unet_ratio, "unet")
                all_params.extend(params)
                lr_descriptions.extend(descriptions)

            return all_params, lr_descriptions

        network.prepare_optimizer_params = types.MethodType(patched_prepare_optimizer_params, network)
        logger.info("Enabled network_reg_lrs override for LyCORIS optimizer param groups.")

    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 get_text_encoder_outputs_caching_strategy(self, args):
        if args.cache_text_encoder_outputs:
            return strategy_anima.AnimaTextEncoderOutputsCachingStrategy(
                args.cache_text_encoder_outputs_to_disk, args.text_encoder_batch_size, args.skip_cache_check, False
            )
        return None

    def cache_text_encoder_outputs_if_needed(
        self, args, accelerator: Accelerator, unet, vae, text_encoders, dataset: train_util.DatasetGroup, weight_dtype
    ):
        if args.cache_text_encoder_outputs:
            if not args.lowram:
                # We cannot move DiT to CPU because of block swap, so only move VAE
                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)

            with accelerator.autocast():
                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 = train_util.load_prompts(args.sample_prompts)
                sample_prompts_te_outputs = {}
                with accelerator.autocast(), 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_and_masks = tokenize_strategy.tokenize(p)
                                sample_prompts_te_outputs[p] = text_encoding_strategy.encode_tokens(
                                    tokenize_strategy, text_encoders, tokens_and_masks
                                )
                self.sample_prompts_te_outputs = sample_prompts_te_outputs

            accelerator.wait_for_everyone()

            # move text encoder back to cpu
            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)
        qwen3_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()
        anima_train_utils.sample_images(
            accelerator,
            args,
            epoch,
            global_step,
            unet,
            vae,
            qwen3_te,
            tokenize_strategy,
            text_encoding_strategy,
            self.sample_prompts_te_outputs,
        )

    def get_noise_scheduler(self, args: argparse.Namespace, device: torch.device) -> Any:
        noise_scheduler = sd3_train_utils.FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.discrete_flow_shift)
        return noise_scheduler

    def encode_images_to_latents(self, args, vae, images):
        vae: qwen_image_autoencoder_kl.AutoencoderKLQwenImage
        return vae.encode_pixels_to_latents(images)  # Keep 4D for input/output

    def shift_scale_latents(self, args, latents):
        # Latents already normalized by vae.encode with scale
        return latents

    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,
    ):
        anima: anima_models.Anima = unet

        # Sample noise
        if latents.ndim == 5:  # Fallback for 5D latents (old cache)
            latents = latents.squeeze(2)  # [B, C, 1, H, W] -> [B, C, H, W]
        batch_size = latents.shape[0]

        if self._random_noise_shift_current is None:
            self._random_noise_shift_current = float(args.random_noise_shift)
        if self._random_noise_multiplier_current is None:
            self._random_noise_multiplier_current = float(args.random_noise_multiplier)

        random_noise_shift_base = self._random_noise_shift_current
        random_noise_multiplier_base = self._random_noise_multiplier_current

        if args.knn_noise_k > 0:
            # For KNN mode, draw K candidates and apply one shared random shift/multiplier per sample across K,
            # then select nearest by distance.
            candidates = torch.randn(
                (batch_size, args.knn_noise_k, *latents.shape[1:]), device=latents.device, dtype=latents.dtype
            )

            if random_noise_shift_base > 0.0:
                if args.random_noise_shift_random_strength:
                    random_noise_shift = torch.rand(1, device=latents.device, dtype=latents.dtype) * random_noise_shift_base
                else:
                    random_noise_shift = random_noise_shift_base
                shared_shift = (
                    torch.randn(batch_size, latents.shape[1], 1, 1, device=latents.device, dtype=latents.dtype)
                    * random_noise_shift
                )
                candidates = candidates + shared_shift.unsqueeze(1)

            if random_noise_multiplier_base > 0.0:
                if args.random_noise_multiplier_random_strength:
                    random_noise_multiplier = (
                        torch.rand(1, device=latents.device, dtype=latents.dtype) * random_noise_multiplier_base
                    )
                else:
                    random_noise_multiplier = random_noise_multiplier_base
                shared_multiplier = torch.exp(
                    torch.randn(batch_size, 1, 1, 1, device=latents.device, dtype=latents.dtype) * random_noise_multiplier
                )
                candidates = candidates * shared_multiplier.unsqueeze(1)

            noise = train_util.select_nearest_noise_candidate(latents, candidates)
        else:
            noise = train_util.sample_training_noise(args, latents)

            if random_noise_shift_base > 0.0:
                if args.random_noise_shift_random_strength:
                    random_noise_shift = torch.rand(1, device=noise.device, dtype=noise.dtype) * random_noise_shift_base
                else:
                    random_noise_shift = random_noise_shift_base
                noise_shift = torch.randn(
                    batch_size,
                    latents.shape[1],
                    1,
                    1,
                    device=noise.device,
                    dtype=noise.dtype,
                ) * random_noise_shift
                noise = noise + noise_shift

            if random_noise_multiplier_base > 0.0:
                if args.random_noise_multiplier_random_strength:
                    random_noise_multiplier = (
                        torch.rand(1, device=noise.device, dtype=noise.dtype) * random_noise_multiplier_base
                    )
                else:
                    random_noise_multiplier = random_noise_multiplier_base
                noise_multiplier = torch.exp(
                    torch.randn(batch_size, 1, 1, 1, device=noise.device, dtype=noise.dtype) * random_noise_multiplier
                )
                noise = noise * noise_multiplier

        if is_train:
            self._random_noise_shift_current = self._random_noise_shift_current * args.random_noise_shift_decay
            self._random_noise_multiplier_current = self._random_noise_multiplier_current * args.random_noise_multiplier_decay

        # Get noisy model input and timesteps
        noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timesteps(
            args, noise_scheduler, latents, noise, accelerator.device, weight_dtype
        )
        timesteps = timesteps / 1000.0  # scale to [0, 1] range. timesteps is float32

        # Gradient checkpointing support
        if args.gradient_checkpointing:
            noisy_model_input.requires_grad_(True)
            for t in text_encoder_conds:
                if t is not None and t.dtype.is_floating_point:
                    t.requires_grad_(True)

        # Unpack text encoder conditions
        prompt_embeds, attn_mask, t5_input_ids, t5_attn_mask = text_encoder_conds

        # Move to device
        prompt_embeds = prompt_embeds.to(accelerator.device, dtype=weight_dtype)
        attn_mask = attn_mask.to(accelerator.device)
        t5_input_ids = t5_input_ids.to(accelerator.device, dtype=torch.long)
        t5_attn_mask = t5_attn_mask.to(accelerator.device)

        # Create padding mask
        bs = latents.shape[0]
        h_latent = latents.shape[-2]
        w_latent = latents.shape[-1]
        padding_mask = torch.zeros(bs, 1, h_latent, w_latent, dtype=weight_dtype, device=accelerator.device)

        # Call model
        noisy_model_input = noisy_model_input.unsqueeze(2)  # 4D to 5D, [B, C, H, W] -> [B, C, 1, H, W]
        with torch.set_grad_enabled(is_train), accelerator.autocast():
            model_pred = anima(
                noisy_model_input,
                timesteps,
                prompt_embeds,
                padding_mask=padding_mask,
                target_input_ids=t5_input_ids,
                target_attention_mask=t5_attn_mask,
                source_attention_mask=attn_mask,
            )
        model_pred = model_pred.squeeze(2)  # 5D to 4D, [B, C, 1, H, W] -> [B, C, H, W]

        # Rectified flow target: noise - latents
        target = noise - latents

        # Loss weighting
        weighting = anima_train_utils.compute_loss_weighting_for_anima(weighting_scheme=args.weighting_scheme, sigmas=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:
        """Override base process_batch for caption dropout with cached text encoder outputs."""

        # Text encoder conditions
        text_encoder_outputs_list = batch.get("text_encoder_outputs_list", None)
        anima_text_encoding_strategy: strategy_anima.AnimaTextEncodingStrategy = text_encoding_strategy
        if text_encoder_outputs_list is not None:
            caption_dropout_rates = text_encoder_outputs_list[-1]
            text_encoder_outputs_list = text_encoder_outputs_list[:-1]

            # Apply caption dropout to cached outputs
            text_encoder_outputs_list = anima_text_encoding_strategy.drop_cached_text_encoder_outputs(
                *text_encoder_outputs_list, caption_dropout_rates=caption_dropout_rates
            )
            batch["text_encoder_outputs_list"] = text_encoder_outputs_list

        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 train_util.get_sai_model_spec_dataclass(None, args, False, True, False, anima="preview").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_knn_noise_k"] = args.knn_noise_k
        metadata["ss_random_noise_shift"] = args.random_noise_shift
        metadata["ss_random_noise_multiplier"] = args.random_noise_multiplier
        metadata["ss_random_noise_shift_random_strength"] = args.random_noise_shift_random_strength
        metadata["ss_random_noise_multiplier_random_strength"] = args.random_noise_multiplier_random_strength
        metadata["ss_random_noise_shift_decay"] = args.random_noise_shift_decay
        metadata["ss_random_noise_multiplier_decay"] = args.random_noise_multiplier_decay

    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 first parameter's requires_grad to True to workaround Accelerate gradient checkpointing bug
        first_param = next(text_encoder.parameters())
        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:
            return super().prepare_unet_with_accelerator(args, accelerator, unet)

        model = unet
        model = accelerator.prepare(model, 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()

        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()
    train_util.add_dit_training_arguments(parser)
    anima_train_utils.add_anima_training_arguments(parser)
    parser.add_argument(
        "--random_noise_shift",
        type=float,
        default=0.0,
        help="stddev of per-sample per-channel random noise shift (disabled when 0.0)",
    )
    parser.add_argument(
        "--random_noise_multiplier",
        type=float,
        default=0.0,
        help="stddev of log-normal random noise multiplier (disabled when 0.0)",
    )
    parser.add_argument(
        "--random_noise_shift_random_strength",
        action="store_true",
        help="use random strength between 0~random_noise_shift for random noise shift",
    )
    parser.add_argument(
        "--random_noise_multiplier_random_strength",
        action="store_true",
        help="use random strength between 0~random_noise_multiplier for random noise multiplier",
    )
    parser.add_argument(
        "--random_noise_shift_decay",
        type=float,
        default=1.0,
        help="decay factor for random_noise_shift applied every training step (0.0-1.0)",
    )
    parser.add_argument(
        "--random_noise_multiplier_decay",
        type=float,
        default=1.0,
        help="decay factor for random_noise_multiplier applied every training step (0.0-1.0)",
    )
    # parser.add_argument("--fp8_scaled", action="store_true", help="Use scaled fp8 for DiT / DiTにスケーリングされたfp8を使う")
    parser.add_argument(
        "--unsloth_offload_checkpointing",
        action="store_true",
        help="offload activations to CPU RAM using async non-blocking transfers (faster than --cpu_offload_checkpointing). "
        "Cannot be used with --cpu_offload_checkpointing or --blocks_to_swap.",
    )
    return parser


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

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

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

    trainer = AnimaNetworkTrainer()
    trainer.train(args)