"""UltraTex multi-view albedo LoRA training on Krea 2 or Qwen-Image 2.1. Same recipe as ``train_flux2.py`` (TexVerse 2x3 canvases, background token dropping, foreground-restricted flow matching, rank-64 LoRA, UNO shifted timestep schedule, periodic LoRA checkpoints with validation renders), with the backbone-specific pieces delegated to ``ultratex.mv``: --backbone krea2 Krea 2 Raw (12B single-stream DiT, 16ch 8x latent, 2x2 patch) --backbone qwen21 Qwen-Image 2.1 (7B single-stream DiT, 64ch 16x RGBA latent, patch 1, block-causal conditioning with cached text/reference prefix) Text conditioning comes from a cache built by ``tools/precompute_text_conditioning.py`` so no text encoder is loaded at train time. """ from __future__ import annotations import dataclasses import gc import json import logging import os import re import time from typing import Literal import torch import transformers from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import set_seed from diffusers.optimization import get_scheduler from peft import LoraConfig, get_peft_model from safetensors.torch import load_file, save_file from torch.utils.data import DataLoader from tqdm import tqdm from ultratex import mv from ultratex.backbones.common import configure_sparse_attention, gather_sequence from ultratex.data.dataset import FluxPairedDatasetV2, LengthAwareBatchSampler logger = get_logger(__name__) @dataclasses.dataclass class TrainArgs: backbone: Literal["krea2", "qwen21"] = "qwen21" model_path: str | None = None # defaults to $KREA2_MODEL_PATH / $QWEN21_MODEL_PATH vae_path: str | None = None # defaults to $KREA2_VAE_PATH / $QWEN21_VAE_PATH text_cond_path: str = "" # tools/precompute_text_conditioning.py output decoder_ckpt: str | None = None # foreground-aware decoder (train_vae/train_decoder_mv.py); optional comfy_root: str | None = None # defaults to $COMFY_ROOT debug_model: bool = False # 1-block random model, for pipeline smoke tests only mixed_precision: Literal["no", "fp16", "bf16"] = "bf16" gradient_accumulation_steps: int = 1 seed: int = 42 lora_rank: int = 64 lora_alpha: int | None = None block_indices: list[int] | None = dataclasses.field(default=None) gradient_checkpoint: bool = True attention: str | None = None # krea2: full; qwen21: block_causal (default per backbone) ref_timestep_zero: bool = False # krea2 only: modulate reference tokens from t = 0 sparse_attention: bool = True sparse_attention_topk: float = 0.20 sparse_attention_blkq: int = 128 sparse_attention_blkk: int = 64 krea2_mu: float = 1.15 # Krea 2 flow shift (mu); Qwen-Image 2.1 derives mu from the token count background: str | None = None # canonical background colour for compose/decode (per-backbone default) learning_rate: float = 1e-4 adam_betas: list[float] = dataclasses.field(default_factory=lambda: [0.9, 0.999]) adam_eps: float = 1e-8 adam_weight_decay: float = 0.01 max_grad_norm: float = 1.0 lr_scheduler: str = "constant" lr_warmup_steps: int = 0 max_train_steps: int = 1000000 train_data_json: str = "data/train_demo.json" eval_data_json: str = "data/eval_demo.json" data_root: str | None = None batch_size: int = 1 resolution: int = 2048 num_workers: int = 8 bucket_metadata_json: str | None = None degrade_single_render: bool = False degrade_single_render_resolution: int = 512 inference_num_steps: int = 25 inference_drop_background_tokens: bool = True inference_remove_background: bool = True eval_max_samples: int | None = None skip_eval: bool = False resume_from_checkpoint: str | None = None checkpointing_steps: int = 2000 save_only_steps: int = 1000 project_dir: str | None = None log_every: int = 10 def __post_init__(self): if self.project_dir is None: self.project_dir = f"outputs/train_{self.backbone}_albedo_res={self.resolution}_lr={self.learning_rate}" if self.lora_alpha is None: self.lora_alpha = self.lora_rank def apply_lora(model, args: TrainArgs): model.requires_grad_(False) config = LoraConfig( r=args.lora_rank, lora_alpha=args.lora_alpha, target_modules=mv.lora_targets(args.backbone, model, args.block_indices), lora_dropout=0.0, bias="none", ) model = get_peft_model(model, config) model.print_trainable_parameters() return model def save_lora(model, accelerator: Accelerator, save_path: str, args: TrainArgs) -> None: os.makedirs(save_path, exist_ok=True) unwrapped = accelerator.unwrap_model(model) state = {k: v.detach().to(torch.float32).cpu().contiguous() for k, v in unwrapped.state_dict().items() if "lora_" in k} save_file(state, os.path.join(save_path, "dit_lora.safetensors"), metadata={"backbone": args.backbone, "lora_rank": str(args.lora_rank)}) with open(os.path.join(save_path, "train_args.json"), "wt", encoding="utf-8") as f: json.dump(dataclasses.asdict(args), f, indent=2) def load_lora(model, path: str) -> int: state = load_file(os.path.join(path, "dit_lora.safetensors")) result = model.load_state_dict(state, strict=False) unexpected = [k for k in result.unexpected_keys if "lora_" in k] if unexpected: raise ValueError(f"LoRA checkpoint {path} does not match the model: unexpected keys {unexpected[:5]}") m = re.fullmatch(r"checkpoint-(\d+)", os.path.basename(os.path.normpath(path))) return int(m.group(1)) if m else 0 def snapshot_lora_base_weights(model): snapshots = [] for name, module in model.named_modules(): if hasattr(module, "lora_A") and hasattr(module, "get_base_layer"): weight = module.get_base_layer().weight snapshots.append((name, weight, weight.detach().to(device="cpu", copy=True))) return snapshots @torch.no_grad() def restore_lora_base_weights(snapshots): for _name, weight, original in snapshots: weight.copy_(original) def run_validation(pipe: mv.MVPipeline, model, eval_dataloader, save_path: str, args: TrainArgs, background_latent, accelerator): for i, batch in enumerate(eval_dataloader): if args.eval_max_samples is not None and i >= args.eval_max_samples: break sample_id = batch["id"][0] output_path = os.path.join(save_path, f"eval_{sample_id}_{i}.webp") seq = pipe.build(batch, drop_background=args.inference_drop_background_tokens) generated = pipe.sample(model, seq, args.inference_num_steps, args.resolution, seed=0, krea2_mu=args.krea2_mu) decoded = pipe.compose_and_decode(seq, generated, background_latent) mask_full = batch["masks_full"][0][0] if args.inference_remove_background else None result = mv.tensor_to_pil(decoded, mask_full, whiten_background=args.inference_remove_background) result.save(os.path.join(save_path, f"result_{sample_id}_{i}.webp"), format="WEBP", lossless=True) composite = mv.make_eval_composite(batch["image_dir"][0], result, batch["render_path"][0], args.resolution) composite.save(output_path, format="WEBP", quality=92) logger.info("validation %s -> %s", sample_id, output_path) def main(args: TrainArgs) -> None: accelerator = Accelerator( project_dir=args.project_dir, gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, ) set_seed(args.seed, device_specific=True) weight_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "no": torch.float32}[accelerator.mixed_precision] logging.basicConfig(format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", level=logging.INFO, force=True) os.makedirs(args.project_dir, exist_ok=True) logger.info("UltraTex %s training: %s", args.backbone, dataclasses.asdict(args)) spec = mv.get_spec(args.backbone) device = accelerator.device comfy_root = args.comfy_root or os.environ.get("COMFY_ROOT") vae = mv.load_vae(args.backbone, args.vae_path, device=device, dtype=weight_dtype, comfy_root=comfy_root) if args.decoder_ckpt: vae.load_decoder_state_dict(load_file(args.decoder_ckpt)) logger.info("loaded foreground-aware decoder from %s", args.decoder_ckpt) vae.requires_grad_(False).eval() text_cache = mv.TextCache(args.text_cond_path, args.backbone, device) pipe = mv.MVPipeline(spec, vae, text_cache, device, weight_dtype, attention=args.attention, ref_timestep_zero=args.ref_timestep_zero) model = mv.load_backbone(args.backbone, args.model_path, device="cpu", dtype=weight_dtype, debug_mode=args.debug_model) sparse_config = configure_sparse_attention( enabled=args.sparse_attention, head_dim=model.head_dim, topk=args.sparse_attention_topk, blkq=args.sparse_attention_blkq, blkk=args.sparse_attention_blkk, ) logger.info("attention: mode=%s backend=%s", pipe.attention, sparse_config) model = apply_lora(model, args) model.base_model.model.gradient_checkpointing = args.gradient_checkpoint model.to(device) model.train() optimizer = torch.optim.AdamW( [p for p in model.parameters() if p.requires_grad], lr=args.learning_rate, betas=tuple(args.adam_betas), weight_decay=args.adam_weight_decay, eps=args.adam_eps, ) lr_scheduler = get_scheduler(args.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps) global_step = 0 if args.resume_from_checkpoint: global_step = load_lora(accelerator.unwrap_model(model), args.resume_from_checkpoint) logger.info("resumed LoRA from %s at step %d", args.resume_from_checkpoint, global_step) start_step = global_step dataset = FluxPairedDatasetV2( json_file=args.train_data_json, resolution=args.resolution, data_root=args.data_root, bucket_metadata_json=args.bucket_metadata_json or None, degrade_single_render_enabled=args.degrade_single_render, degrade_single_render_resolution=args.degrade_single_render_resolution, ) sampler = LengthAwareBatchSampler(dataset, batch_size=args.batch_size, drop_last=False) dataloader = DataLoader(dataset, batch_sampler=sampler, collate_fn=dataset.collate_fn, num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0) eval_dataset = FluxPairedDatasetV2(json_file=args.eval_data_json, resolution=args.resolution, data_root=args.data_root, degrade_single_render_enabled=args.degrade_single_render, degrade_single_render_resolution=args.degrade_single_render_resolution) eval_dataloader = DataLoader(eval_dataset, batch_size=1, shuffle=False, collate_fn=eval_dataset.collate_fn) logger.info("train objects=%d eval objects=%d", len(dataset.data_dicts), len(eval_dataset.data_dicts)) model, optimizer, lr_scheduler, dataloader = accelerator.prepare(model, optimizer, lr_scheduler, dataloader) background = args.background or spec.default_background background_latent = mv.canonical_background_latent(vae, spec, args.resolution, background) training_timesteps = pipe.training_timesteps(args.resolution, args.krea2_mu) logger.info("timestep schedule: mu=%.3f background=%s", pipe.mu(args.resolution, args.krea2_mu), background) log_path = os.path.join(args.project_dir, "train_log.jsonl") progress = tqdm(range(args.max_train_steps), initial=start_step, total=args.max_train_steps, desc="Steps", disable=not accelerator.is_local_main_process) train_loss, window = 0.0, 0 last_time = time.time() while global_step < args.max_train_steps: for batch in dataloader: seq = pipe.build(batch, drop_background=True) if global_step == start_step: logger.info("BTD: target %d/%d tokens kept, reference tokens kept %s", seq.target_tokens.shape[1], seq.target_token_count, [r.shape[1] for r in seq.ref_tokens]) x_1 = seq.target_tokens.float() t = training_timesteps[torch.randint(0, training_timesteps.shape[0], (x_1.shape[0],), device=device)] x_0 = torch.randn_like(x_1) x_t = (1 - t[:, None, None]) * x_1 + t[:, None, None] * x_0 with accelerator.accumulate(model): pred = pipe.predict(model, seq, x_t, t) loss = mv.flow_loss(pred, x_0, x_1) accelerator.backward(loss) if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) if not torch.isfinite(torch.as_tensor(grad_norm)).all(): optimizer.zero_grad() raise FloatingPointError(f"non-finite gradient norm before step {global_step + 1} (ids={batch.get('id')})") optimizer.step() lr_scheduler.step() optimizer.zero_grad() if accelerator.sync_gradients: global_step += 1 progress.update(1) train_loss += loss.detach().item() window += 1 if global_step % args.log_every == 0: now = time.time() record = {"step": global_step, "loss": train_loss / window, "lr": lr_scheduler.get_last_lr()[0], "sec_per_step": (now - last_time) / window, "tokens": int(seq.target_tokens.shape[1] + sum(r.shape[1] for r in seq.ref_tokens)), "max_mem_gib": torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else None} with open(log_path, "at", encoding="utf-8") as f: f.write(json.dumps(record) + "\n") progress.set_postfix(loss=f"{record['loss']:.4f}", spp=f"{record['sec_per_step']:.1f}") train_loss, window, last_time = 0.0, 0, now should_checkpoint = global_step % args.checkpointing_steps == 0 or global_step == start_step + 100 should_save_only = global_step % args.save_only_steps == 0 and not should_checkpoint if should_save_only or should_checkpoint: save_path = os.path.join(args.project_dir, f"checkpoint-{global_step}") if accelerator.is_main_process: save_lora(model, accelerator, save_path, args) logger.info("saved LoRA to %s", save_path) if should_checkpoint and not args.skip_eval and accelerator.is_main_process: model.eval() inference_model = accelerator.unwrap_model(model) snapshots = snapshot_lora_base_weights(inference_model) inference_model.merge_adapter(safe_merge=True) torch.cuda.empty_cache() try: with torch.no_grad(): run_validation(pipe, inference_model, eval_dataloader, save_path, args, background_latent, accelerator) finally: inference_model.unmerge_adapter() restore_lora_base_weights(snapshots) del snapshots torch.cuda.empty_cache() gc.collect() model.train() if global_step >= args.max_train_steps: break accelerator.wait_for_everyone() accelerator.end_training() if __name__ == "__main__": parser = transformers.HfArgumentParser([TrainArgs]) (args,) = parser.parse_args_into_dataclasses(args_file_flag="--config") main(args)