| """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 |
| vae_path: str | None = None |
| text_cond_path: str = "" |
| decoder_ckpt: str | None = None |
| comfy_root: str | None = None |
| debug_model: bool = False |
|
|
| 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 |
| ref_timestep_zero: bool = False |
| 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 |
| background: str | None = None |
|
|
| 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) |
|
|
