ArtiAgent - DefectFill/engine/DefectFill/train.py

import os
import json
import torch
import argparse
import torch.nn.functional as F
import random
import time
from torch.optim import AdamW
from tqdm import tqdm
from diffusers import DDPMScheduler
from model import DefectFillModel, USE_MODELSCOPE
from data_loader import get_data_loaders
from utils import save_checkpoint, load_checkpoint
# TensorBoard support
from torch.utils.tensorboard import SummaryWriter
import datetime


def generate_seed_from_timestamp():
    """Generates a random seed based on the current timestamp"""
    return int(time.time() * 1000) % (2**31)

def train(args):
    # Set up device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    if torch.cuda.is_available():
        print(f"GPU Name: {torch.cuda.get_device_name(0)}")
    else:
        print("CUDA is NOT available. The model will run on CPU, causing severe slowdown.")
    
    # Create output directory structure
    os.makedirs(args.output_dir, exist_ok=True)
    checkpoints_dir = os.path.join(args.output_dir, "checkpoints")
    tensorboard_dir = os.path.join(args.output_dir, "tensorboard")
    os.makedirs(checkpoints_dir, exist_ok=True)
    os.makedirs(tensorboard_dir, exist_ok=True)
    
    # Create log file in the output directory
    log_file_path = os.path.join(args.output_dir, "train_log.txt")
    log_file = open(log_file_path, "a")
    log_file.write(f"\n\n{'='*60}\n")
    log_file.write(f"Training started at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
    log_file.write(f"Object class: {args.object_class}\n")
    log_file.write(f"Defect type: {args.defect_type if args.defect_type else 'all'}\n")
    log_file.write(f"Config name: {args.config_name}\n")
    log_file.write(f"Lambda defect: {args.lambda_defect}\n")
    log_file.write(f"Lambda obj: {args.lambda_obj}\n")
    log_file.write(f"Lambda attn: {args.lambda_attn}\n")
    log_file.write(f"Alpha (obj branch bg weight): {args.alpha}\n")
    log_file.write(f"Gradient accumulation steps: {args.gradient_accumulation_steps}\n")
    log_file.write(f"Random seed: {args.seed}\n")
    log_file.write(f"{'='*60}\n\n")
    
    # Save training configuration to JSON
    train_config = {
        "timestamp": datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'),
        "object_class": args.object_class,
        "defect_type": args.defect_type if args.defect_type else "all",
        "config_name": args.config_name,
        "lambda_defect": args.lambda_defect,
        "lambda_obj": args.lambda_obj,
        "lambda_attn": args.lambda_attn,
        "alpha": args.alpha,
        "batch_size": args.batch_size,
        "max_train_steps": args.max_train_steps,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "lora_rank": args.lora_rank,
        "lora_alpha": args.lora_alpha,
        "text_encoder_lr": args.text_encoder_lr,
        "unet_lr": args.unet_lr,
        "lr_warmup_steps": args.lr_warmup_steps,
        "save_steps": args.save_steps,
        "seed": args.seed
    }
    config_path = os.path.join(args.output_dir, "train_config.json")
    with open(config_path, "w") as f:
        json.dump(train_config, f, indent=4)
    print(f"Training config saved to: {config_path}")
    
    # Initialize TensorBoard writer
    writer = SummaryWriter(tensorboard_dir)
    
    # Load data
    train_loader, test_loader = get_data_loaders(
        root_dir=args.data_dir,
        object_class=args.object_class,
        batch_size=args.batch_size,
        defect_type=args.defect_type,
        dilate_mask=args.dilate_mask,
        mask_kernel_size=args.mask_kernel_size
    )
    
    # Initialize model
    model = DefectFillModel(
        device=device,
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
        seed=args.seed
    )
    
    # Set up optimizers with specific learning rates for Text Encoder and UNet
    text_encoder_params = [p for n, p in model.pipeline.text_encoder.named_parameters() if "lora" in n]
    unet_params = [p for n, p in model.pipeline.unet.named_parameters() if "lora" in n]
    
    optimizer = AdamW([
        {"params": text_encoder_params, "lr": args.text_encoder_lr},
        {"params": unet_params, "lr": args.unet_lr}
    ])
    
    # Save original LRs for warmup calculation
    base_lrs = [args.text_encoder_lr, args.unet_lr]
    
    # Set up noise scheduler (handling ModelScope vs HuggingFace)
    hf_model_id = "sd2-community/stable-diffusion-2-inpainting"
    if USE_MODELSCOPE:
        try:
            from modelscope import snapshot_download
            print(f"[ModelScope] Downloading scheduler: {hf_model_id}")
            local_model_path = snapshot_download(hf_model_id)
            noise_scheduler = DDPMScheduler.from_pretrained(local_model_path, subfolder="scheduler")
        except ImportError:
            print("[Warning] modelscope not installed, falling back to HuggingFace...")
            noise_scheduler = DDPMScheduler.from_pretrained(hf_model_id, subfolder="scheduler")
    else:
        noise_scheduler = DDPMScheduler.from_pretrained(hf_model_id, subfolder="scheduler")
    
    # Resume from checkpoint if specified
    start_step = 0
    if args.resume_from:
        start_step = load_checkpoint(model, optimizer, args.resume_from)
        print(f"Resuming from step {start_step}")
        log_file.write(f"Resuming from step {start_step}\n")
    
    # Set models to training mode
    model.pipeline.unet.train()
    model.pipeline.text_encoder.train()
    
    total_steps = args.max_train_steps
    progress_bar = tqdm(range(start_step, total_steps), desc="Training Progress")
    
    global_step = start_step
    accumulation_step = 0

    # Add this BEFORE the while loop
    with torch.no_grad():
        clean_latents_cache = {}
        for batch in train_loader:
            images = batch["image"].to(device, dtype=torch.float16)
            is_defect = batch["is_defect"]
            defect_samples = torch.nonzero(is_defect).squeeze(1)
            if len(defect_samples) > 0:
                defect_images = images[defect_samples]
                # Cache original latents
                latents = model.pipeline.vae.encode(defect_images).latent_dist.sample() * model.pipeline.vae.config.scaling_factor
                clean_latents_cache[tuple(defect_samples.cpu().tolist())] = latents.detach()
    
    while global_step < total_steps:
        for batch in train_loader:
            if global_step >= total_steps:
                break

            # t_start = time.time()

            # Move data to device
            images = batch["image"].to(device, dtype=torch.float16)
            masks = batch["mask"].to(device, dtype=torch.float16)
            backgrounds = batch["background"].to(device, dtype=torch.float16)
            adjusted_masks = batch["adjusted_mask"].to(device, dtype=torch.float16)
            is_defect = batch["is_defect"]
            
            # Ensure we only process defective samples
            defect_samples = torch.nonzero(is_defect).squeeze(1)
            if len(defect_samples) == 0:
                continue  # Skip batch if no defects present
                
            # Extract defect-only samples
            defect_images = images[defect_samples]
            defect_masks = masks[defect_samples]
            defect_backgrounds = backgrounds[defect_samples]
            defect_adjusted_masks = adjusted_masks[defect_samples]
            object_classes = [batch["object_class"][i] for i in defect_samples]

            # t_data = time.time()
            
            # Extract defect type from file path for specific prompting
            defect_types = []
            for i in defect_samples:
                if hasattr(train_loader.dataset, 'images') and i < len(train_loader.dataset.images):
                    img_path = train_loader.dataset.images[i]
                    parts = img_path.split(os.sep)
                    for j, part in enumerate(parts):
                        if part == "defective" and j + 1 < len(parts):
                            defect_types.append(parts[j + 1])
                            break
                    else:
                        defect_types.append("defect")
                else:
                    defect_types.append("defect")
            
            # Learning rate warmup
            if global_step < args.lr_warmup_steps:
                lr_scale = min(1.0, (global_step + 1) / args.lr_warmup_steps)
                for i, param_group in enumerate(optimizer.param_groups):
                    param_group["lr"] = base_lrs[i] * lr_scale
            
            # Reset attention maps
            if hasattr(model, 'attention_maps'):
                model.attention_maps = {}
            
            # ========== PHASE 1: Defect Branch (Defect Texture Learning) ==========
            # Using <defect> learnable token for concept isolation
            # t_text = time.time()
            defect_prompts = [f"A photo of {model.placeholder_token}" for _ in range(len(defect_samples))]
            text_embeddings = model.get_text_embeddings(defect_prompts, enable_grad=True)
            
            # ========== PHASE 1 & 2 OPTIMIZED VAE & MASK PASS ==========
            
            # 1. Single VAE encoding for original defect images
            with torch.no_grad():
                # Instead of re-encoding, load from cache
                latents = clean_latents_cache[tuple(defect_samples.cpu().tolist())]

            # t_vae = time.time()
            
            if len(defect_masks.shape) == 3:
                defect_masks = defect_masks.unsqueeze(1)

            # 2. Fully Vectorized Random Mask Generation directly on GPU (0 Python Loops)
            B, _, H, W = defect_images.shape
            K = 15  # Number of random boxes

            # Create broadcastable coordinate grids (1, 1, H, 1) and (1, 1, 1, W)
            grid_y = torch.arange(H, device=device).view(1, 1, H, 1)
            grid_x = torch.arange(W, device=device).view(1, 1, 1, W)

            # Sample all box parameters for all batch items and box counts simultaneously: shape (B, K, 1, 1)
            rh = torch.randint(max(1, int(H * 0.03)), max(2, int(H * 0.25)), (B, K, 1, 1), device=device)
            rw = torch.randint(max(1, int(W * 0.03)), max(2, int(W * 0.25)), (B, K, 1, 1), device=device)
            ry = torch.randint(0, max(1, H - 20), (B, K, 1, 1), device=device)
            rx = torch.randint(0, max(1, W - 20), (B, K, 1, 1), device=device)

            # Compute rectangle coverage across all (B, K, H, W) elements in a single broadcasted operation
            in_box = (grid_y >= ry) & (grid_y < (ry + rh)) & (grid_x >= rx) & (grid_x < (rx + rw))

            # Merge all K boxes per image using logical OR (.any), shape -> (B, 1, H, W)
            random_masks = in_box.any(dim=1, keepdim=True).to(dtype=defect_images.dtype)

            # 3. Batched VAE Encoding (Combines Phase 1 and Phase 2 into 1 pass)
            with torch.no_grad():
                masked_images = defect_images * (1 - defect_masks)
                rand_masked_images = defect_images * (1 - random_masks)
                
                combined_masked = torch.cat([masked_images, rand_masked_images], dim=0)
                combined_latents = model.pipeline.vae.encode(combined_masked).latent_dist.sample() * model.pipeline.vae.config.scaling_factor
                
                masked_image_latents, random_masked_image_latents = torch.chunk(combined_latents, 2, dim=0)

            # 4. Latent interpolation for masks
            mask_latents = F.interpolate(defect_masks, size=(latents.shape[2], latents.shape[3]))
            random_mask_latents = F.interpolate(random_masks, size=(latents.shape[2], latents.shape[3]))

            # ========== Phase 1 Forward Pass ==========
            noise = torch.randn_like(latents)
            timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device)
            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

            outputs = model(
                noisy_latents=noisy_latents,
                masked_image_latents=masked_image_latents,
                mask_latents=mask_latents,
                timesteps=timesteps,
                encoder_hidden_states=text_embeddings
            )

            # t_forward = time.time()
            
            # print(f"[TIMING RESULTS]")
            # print(f" ├─ Data Fetch & Filter : {t_data - t_start:.2f}s")
            # print(f" ├─ Text Encoding       : {t_text - t_data:.2f}s")
            # print(f" ├─ VAE Encoding        : {t_vae - t_text:.2f}s")
            # print(f" └─ UNet Forward & Loss : {t_forward - t_vae:.2f}s")

            
            noise_pred = outputs["noise_pred"]
            defect_loss = model.compute_defect_loss(noise_pred, noise, mask_latents)
            attention_loss = outputs.get("attention_loss", torch.tensor(0.0, device=device))

            # ========== Phase 2 Forward Pass ==========
            obj_prompts = [f"A {obj_class} with {model.placeholder_token}" for obj_class in object_classes]
            obj_text_embeddings = model.get_text_embeddings(obj_prompts, enable_grad=True)

            obj_noise = torch.randn_like(latents)
            obj_timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device)
            obj_noisy_latents = noise_scheduler.add_noise(latents, obj_noise, obj_timesteps)
            
            obj_outputs = model(
                noisy_latents=obj_noisy_latents,
                masked_image_latents=random_masked_image_latents,
                mask_latents=random_mask_latents,
                timesteps=obj_timesteps,
                encoder_hidden_states=obj_text_embeddings
            )
            
            object_loss = model.compute_object_loss(obj_outputs["noise_pred"], obj_noise, random_mask_latents, alpha=args.alpha)
            
            # Total Loss Calculation
            total_loss = args.lambda_defect * defect_loss + args.lambda_obj * object_loss + args.lambda_attn * attention_loss
            total_loss = total_loss / args.gradient_accumulation_steps
            
            # NaN Check
            if torch.isnan(total_loss):
                print(f"Warning: NaN loss detected at step {global_step}")
                log_file.write(f"Warning: NaN loss at step {global_step}\n")
                optimizer.zero_grad()
                continue

            # t_before_backward = time.time()
            total_loss.backward()
            # t_after_backward = time.time()
            # print(f" └─ Backward Pass Time : {t_after_backward - t_before_backward:.2f}s")
            accumulation_step += 1
            
            # Optimization step
            if accumulation_step >= args.gradient_accumulation_steps:
                optimizer.step()
                optimizer.zero_grad()
                accumulation_step = 0
                
                progress_bar.update(1)
                global_step += 1
                
                # Logging and TensorBoard updates
                writer.add_scalar("Loss/Defect", defect_loss.item(), global_step)
                writer.add_scalar("Loss/Object", object_loss.item(), global_step)
                writer.add_scalar("Loss/Attention", attention_loss.item(), global_step)
                writer.add_scalar("Loss/Total", total_loss.item() * args.gradient_accumulation_steps, global_step)
                
                if global_step % 10 == 0:
                    for i, param_group in enumerate(optimizer.param_groups):
                        writer.add_scalar(f"LearningRate/group{i}", param_group["lr"], global_step)
                
                # Periodic checkpointing
                if global_step % args.save_steps == 0 or global_step == total_steps:
                    checkpoint_path = os.path.join(checkpoints_dir, f"checkpoint_{global_step}.pt")
                    save_checkpoint(model, optimizer, global_step, checkpoint_path)
                    log_file.write(f"Checkpoint saved at step {global_step}\n")

    # Save final model
    final_checkpoint_path = os.path.join(checkpoints_dir, "checkpoint_final.pt")
    save_checkpoint(model, optimizer, global_step, final_checkpoint_path)
    print(f"Final model saved to: {final_checkpoint_path}")
    
    writer.close()
    log_file.close()
    return model

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train DefectFill model")
    
    # Paths
    parser.add_argument("--data_dir", type=str, required=True, help="Path to MVTec AD dataset")
    parser.add_argument("--object_class", type=str, required=True, help="Object class to train on")
    parser.add_argument("--output_dir", type=str, default="./output", help="Directory to save models")
    
    # Loss Weights
    parser.add_argument("--lambda_defect", type=float, default=0.5, help="Defect loss weight (L_def)")
    parser.add_argument("--lambda_obj", type=float, default=0.2, help="Object integrity loss weight (L_obj)")
    parser.add_argument("--lambda_attn", type=float, default=0.05, help="Attention loss weight (L_attn)")
    parser.add_argument("--alpha", type=float, default=0.3, help="Background weight for object branch")
    
    # Training Config
    parser.add_argument("--config_name", type=str, default="base", help="Experiment name (base/tex/obj)")
    parser.add_argument("--defect_type", type=str, default=None, help="Specific defect type to train on")
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--lora_rank", type=int, default=8)
    parser.add_argument("--lora_alpha", type=int, default=16)
    parser.add_argument("--text_encoder_lr", type=float, default=4e-5)
    parser.add_argument("--unet_lr", type=float, default=2e-4)
    parser.add_argument("--max_train_steps", type=int, default=2000)
    parser.add_argument("--lr_warmup_steps", type=int, default=100)
    parser.add_argument("--save_steps", type=int, default=500)
    parser.add_argument("--seed", type=int, default=-1, help="-1 for timestamp-based seed")
    parser.add_argument("--resume_from", type=str, default=None)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=2)
    parser.add_argument("--dilate_mask", type=str, default="False", help="Whether to dilate masks (True/False)")
    parser.add_argument("--mask_kernel_size", type=int, default=3, help="Size of dilation kernel (must be odd, e.g. 3, 5, 7)")
    
    args = parser.parse_args()

    # Helper to convert string "True" to boolean
    args.dilate_mask = args.dilate_mask.lower() == "true"
    
    # Seed setup
    if args.seed == -1:
        args.seed = generate_seed_from_timestamp()
    
    torch.manual_seed(args.seed)
    random.seed(args.seed)
    
    train(args)