from tqdm import tqdm from model import GameNGen, ActionEncoder import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from config import ModelConfig, TrainingConfig import pandas as pd from torchvision import transforms import os from PIL import Image import json import logging import torch.nn.functional as F from diffusers.optimization import get_cosine_schedule_with_warmup from accelerate import Accelerator from huggingface_hub import hf_hub_download from peft import LoraConfig import mlflow import argparse class NextFrameDataset(Dataset): def __init__(self, num_actions: int, metadata_path: str, frames_dir: str, image_size: tuple, history_len: int, subset_percentage: float): self.metadata = pd.read_csv(metadata_path) self.frames_dir = frames_dir # List files and filter out non-image files if necessary self.frame_files = sorted( [f for f in os.listdir(frames_dir) if f.endswith('.png')], key=lambda x: int(x.split('_')[1].split('.')[0]) ) # Calculate the number of frames to use based on the percentage num_to_use = int(len(self.frame_files) * subset_percentage) self.frame_files = self.frame_files[:num_to_use] self.metadata = self.metadata.iloc[:num_to_use] print(f"Using a {subset_percentage*100}% subset of the data: {len(self.frame_files)} frames.") self.num_actions = num_actions self.total_frames = len(self.frame_files) self.history_len = history_len self.transform = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) # Normalize VAE to [-1, 1] ]) def __len__(self) -> int: # We can't use the first `history_len` frames as they don't have enough history return min(len(self.metadata), self.total_frames) - self.history_len - 1 def __getitem__(self, idx: int) -> dict: # We are getting the item at `idx` in our shortened dataset. # The actual index in the video/metadata is `idx + self.history_len`. actual_idx = idx + self.history_len history_frames = [] for i in range(self.history_len): frame_idx = actual_idx - self.history_len + i # Use the sorted file list to get the correct frame img_path = os.path.join(self.frames_dir, self.frame_files[frame_idx]) try: pil_image = Image.open(img_path).convert("RGB") except FileNotFoundError: raise IndexError(f"Could not read history frame {frame_idx} from {img_path}.") history_frames.append(self.transform(pil_image)) history_tensor = torch.stack(history_frames) # Get the target frame (next_frame) next_frame_img_path = os.path.join(self.frames_dir, self.frame_files[actual_idx]) try: next_pil_image = Image.open(next_frame_img_path).convert("RGB") except FileNotFoundError: raise IndexError(f"Could not read frame {actual_idx} from {next_frame_img_path}.") next_image = self.transform(next_pil_image) # Get the action that led to the `next_frame` action_row = self.metadata.iloc[actual_idx] action_data = json.loads(str(action_row['action'])) action_int = int(action_data[0] if isinstance(action_data, list) else action_data) curr_action = torch.zeros(self.num_actions) curr_action[action_int] = 1.0 return { "frame_history": history_tensor, "action": curr_action, "next_frame": next_image } def train(): logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s") parser = argparse.ArgumentParser(description="GameNGen Finetuning") parser.add_argument("--metadata_input", type=str, required=True, help="Path to the metadata CSV file") parser.add_argument("--frames_input", type=str, required=True, help="Path to the frames directory") parser.add_argument("--experiment_name", type=str, default="GameNGen Finetuning", help="Name of the MLflow experiment.") args = parser.parse_args() # --- MLflow Integration --- # Check for Azure ML environment. # The v1 SDK may set AZUREML_MLFLOW_URI, while v2 sets MLFLOW_TRACKING_URI. is_azureml_env = "AZUREML_MLFLOW_URI" in os.environ or \ ("MLFLOW_TRACKING_URI" in os.environ and "azureml" in os.environ["MLFLOW_TRACKING_URI"]) if is_azureml_env: # In Azure ML, MLflow is configured automatically by environment variables. # We don't need to set the tracking URI or experiment name. logging.info("✅ MLflow using Azure ML environment configuration.") else: # For local runs, explicitly set up a local tracking URI and experiment. # This will save runs to a local 'mlruns' directory. mlflow.set_tracking_uri("file:./mlruns") mlflow.set_experiment(args.experiment_name) logging.info(f"⚠️ Using local MLflow tracking (./mlruns) for experiment '{args.experiment_name}'.") # --- Setup --- accelerator = Accelerator( mixed_precision="fp16", gradient_accumulation_steps=1 ) model_config = ModelConfig() train_config = TrainingConfig() # Define file paths using the config metadata_path = args.metadata_input frames_dir = args.frames_input engine = GameNGen(model_config.model_id, model_config.num_timesteps, history_len=model_config.history_len) # --- Memory Saving Optimizations --- engine.unet.enable_gradient_checkpointing() # try: # engine.unet.enable_xformers_memory_efficient_attention() # logging.info("xformers memory-efficient attention enabled.") # except ImportError: # logging.warning("xformers is not installed. For better memory efficiency, run: pip install xformers") dataset = NextFrameDataset(model_config.num_actions, metadata_path, frames_dir, model_config.image_size, history_len=model_config.history_len, subset_percentage=train_config.subset_percentage) dataloader = DataLoader( dataset=dataset, batch_size=train_config.batch_size, shuffle=True, num_workers=0 ) cross_attention_dim = engine.unet.config.cross_attention_dim action_encoder = ActionEncoder(model_config.num_actions, cross_attention_dim) if model_config.use_lora: engine.unet.requires_grad_(False) lora_config = LoraConfig( r=train_config.lora_rank, lora_alpha=train_config.lora_alpha, target_modules=["to_q", "to_k", "to_v", "to_out.0"], lora_dropout=0.1, bias="lora_only", ) engine.unet.add_adapter(lora_config) lora_layers = filter(lambda p: p.requires_grad, engine.unet.parameters()) params_to_train = list(lora_layers) + list(action_encoder.parameters()) else: params_to_train = list(engine.unet.parameters()) + list(action_encoder.parameters()) optim = torch.optim.AdamW(params=params_to_train, lr=train_config.learning_rate) lr_scheduler = get_cosine_schedule_with_warmup( optimizer=optim, num_warmup_steps=500, num_training_steps=len(dataloader) * train_config.num_epochs ) engine, action_encoder, optim, dataloader, lr_scheduler = accelerator.prepare( engine, action_encoder, optim, dataloader, lr_scheduler ) mlflow.autolog(log_models=False) # --- Add an output directory for checkpoints --- output_dir = "./outputs" os.makedirs(output_dir, exist_ok=True) logging.info("Starting training loop...") mlflow.log_params({ "learning_rate": train_config.learning_rate, "batch_size": train_config.batch_size, "num_epochs": train_config.num_epochs, "use_lora": model_config.use_lora, "lora_rank": train_config.lora_rank if model_config.use_lora else None, "subset_percentage": train_config.subset_percentage }) global_step = 0 for epoch in range(train_config.num_epochs): progress_bar = tqdm(total=len(dataloader), disable=not accelerator.is_local_main_process) progress_bar.set_description(f"Epoch {epoch}") for batch in dataloader: optim.zero_grad() next_frames, actions, frame_history = batch["next_frame"], batch["action"], batch["frame_history"] # Encode into latent space with torch.no_grad(): vae = accelerator.unwrap_model(engine).vae latent_dist = vae.encode(next_frames).latent_dist clean_latents = latent_dist.sample() * vae.config.scaling_factor # Encode history frames bs, hist_len, C, H, W = frame_history.shape frame_history = frame_history.view(bs * hist_len, C, H, W) history_latents = vae.encode(frame_history).latent_dist.sample() _, latent_C, latent_H, latent_W = history_latents.shape history_latents = history_latents.reshape(bs, hist_len * latent_C, latent_H, latent_W) # Add noise to history latents to prevent drift (noise augmentation) noise_level = 0.1 # Start with a small, fixed amount of noise history_noise = torch.randn_like(history_latents) * noise_level corrupted_history_latents = history_latents + history_noise # Conditioning is now only the action action_conditioning = action_encoder(actions) conditioning_batch = action_conditioning.unsqueeze(1) # create random noise noise = torch.randn_like(clean_latents) # pick random timestep. High timstep means more noise timesteps = torch.randint(0, engine.scheduler.config.num_train_timesteps, (clean_latents.shape[0], ), device=clean_latents.device).long() noisy_latents = engine.scheduler.add_noise(clean_latents, noise, timesteps) # Concatenate history latents with noisy latents model_input = torch.cat([noisy_latents, corrupted_history_latents], dim=1) with accelerator.accumulate(engine): noise_pred = engine(model_input, timesteps, conditioning_batch) loss = F.mse_loss(noise_pred, noise) accelerator.backward(loss) accelerator.clip_grad_norm_(engine.unet.parameters(), 1.0) optim.step() lr_scheduler.step() progress_bar.update(1) logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step} # Log metrics to MLflow if global_step % 10 == 0: # Log every 10 steps to avoid too much overhead mlflow.log_metric("loss", logs["loss"], step=global_step) mlflow.log_metric("learning_rate", logs["lr"], step=global_step) progress_bar.set_postfix(**logs) global_step += 1 progress_bar.close() if accelerator.is_main_process: logging.info(f"Epoch {epoch} complete. Saving checkpoint...") # Define a unique directory for this epoch's checkpoint checkpoint_dir = os.path.join(output_dir, f"checkpoint_epoch_{epoch}") # Use accelerator.save_state to save everything accelerator.save_state(checkpoint_dir) logging.info(f"Checkpoint saved to {checkpoint_dir}") # Save models at the end of training if accelerator.is_main_process: unwrapped_unet = accelerator.unwrap_model(engine).unet unwrapped_action_encoder = accelerator.unwrap_model(action_encoder) try: # Log the action encoder mlflow.pytorch.log_model(unwrapped_action_encoder, "action_encoder") logging.info("✅ Action encoder logged to MLflow") # Log the UNet (or its LoRA weights) if model_config.use_lora: from peft import get_peft_model_state_dict import json lora_save_path = "unet_lora_weights" os.makedirs(lora_save_path, exist_ok=True) # Save LoRA weights using PEFT method lora_state_dict = get_peft_model_state_dict(unwrapped_unet) torch.save(lora_state_dict, os.path.join(lora_save_path, "pytorch_lora_weights.bin")) # Save adapter config adapter_config = unwrapped_unet.peft_config with open(os.path.join(lora_save_path, "adapter_config.json"), "w") as f: json.dump(adapter_config, f, indent=2, default=str) mlflow.log_artifacts(lora_save_path, artifact_path="unet_lora") logging.info("✅ LoRA weights logged to MLflow") else: mlflow.pytorch.log_model(unwrapped_unet, "unet") logging.info("✅ UNet logged to MLflow") logging.info(f"✅ Training completed. MLflow Run ID: {mlflow.active_run().info.run_id}") except Exception as e: logging.error(f"❌ Error logging models to MLflow: {e}") # Save models locally as fallback torch.save(unwrapped_action_encoder.state_dict(), os.path.join(output_dir, "action_encoder.pth")) if model_config.use_lora: try: from peft import get_peft_model_state_dict lora_state_dict = get_peft_model_state_dict(unwrapped_unet) torch.save(lora_state_dict, os.path.join(output_dir, "lora_weights.bin")) logging.info("📁 LoRA weights saved locally") except Exception as lora_e: logging.error(f"❌ Error saving LoRA weights: {lora_e}") torch.save(unwrapped_unet.state_dict(), os.path.join(output_dir, "unet_full.pth")) logging.info("📁 Full UNet saved locally as fallback") else: torch.save(unwrapped_unet.state_dict(), os.path.join(output_dir, "unet.pth")) logging.info("📁 Models saved locally as fallback") if __name__ == "__main__": train()