import imageio, os, torch, warnings, torchvision, torchaudio, argparse, json, random, math import time import numpy as np from typing import Optional from peft import LoraConfig, inject_adapter_in_model from PIL import Image import pandas as pd from tqdm import tqdm from accelerate import Accelerator from accelerate.utils import DistributedDataParallelKwargs from transformers import get_wsd_schedule from torchcodec.decoders import AudioDecoder from .cache_shards import ShardedBinIdxWriter, ShardedBinIdxReader, DEFAULT_META_FILENAME from safetensors.torch import load_file as safe_load_file class ImageDataset(torch.utils.data.Dataset): def __init__( self, base_path=None, metadata_path=None, max_pixels=1920*1080, height=None, width=None, height_division_factor=16, width_division_factor=16, data_file_keys=("image",), image_file_extension=("jpg", "jpeg", "png", "webp"), repeat=1, args=None, ): if args is not None: base_path = args.dataset_base_path metadata_path = args.dataset_metadata_path height = args.height width = args.width max_pixels = args.max_pixels data_file_keys = args.data_file_keys.split(",") repeat = args.dataset_repeat self.base_path = base_path self.max_pixels = max_pixels self.height = height self.width = width self.height_division_factor = height_division_factor self.width_division_factor = width_division_factor self.data_file_keys = data_file_keys self.image_file_extension = image_file_extension self.repeat = repeat if height is not None and width is not None: print("Height and width are fixed. Setting `dynamic_resolution` to False.") self.dynamic_resolution = False elif height is None and width is None: print("Height and width are none. Setting `dynamic_resolution` to True.") self.dynamic_resolution = True if metadata_path is None: print("No metadata. Trying to generate it.") metadata = self.generate_metadata(base_path) print(f"{len(metadata)} lines in metadata.") self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] elif metadata_path.endswith(".json"): with open(metadata_path, "r") as f: metadata = json.load(f) self.data = metadata elif metadata_path.endswith(".jsonl"): metadata = [] with open(metadata_path, 'r') as f: for line in tqdm(f): metadata.append(json.loads(line.strip())) self.data = metadata else: metadata = pd.read_csv(metadata_path) self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] def generate_metadata(self, folder): image_list, prompt_list = [], [] file_set = set(os.listdir(folder)) for file_name in file_set: if "." not in file_name: continue file_ext_name = file_name.split(".")[-1].lower() file_base_name = file_name[:-len(file_ext_name)-1] if file_ext_name not in self.image_file_extension: continue prompt_file_name = file_base_name + ".txt" if prompt_file_name not in file_set: continue with open(os.path.join(folder, prompt_file_name), "r", encoding="utf-8") as f: prompt = f.read().strip() image_list.append(file_name) prompt_list.append(prompt) metadata = pd.DataFrame() metadata["image"] = image_list metadata["prompt"] = prompt_list return metadata def crop_and_resize(self, image, target_height, target_width): width, height = image.size scale = max(target_width / width, target_height / height) image = torchvision.transforms.functional.resize( image, (round(height*scale), round(width*scale)), interpolation=torchvision.transforms.InterpolationMode.BILINEAR ) image = torchvision.transforms.functional.center_crop(image, (target_height, target_width)) return image def get_height_width(self, image): if self.dynamic_resolution: width, height = image.size if width * height > self.max_pixels: scale = (width * height / self.max_pixels) ** 0.5 height, width = int(height / scale), int(width / scale) height = height // self.height_division_factor * self.height_division_factor width = width // self.width_division_factor * self.width_division_factor else: height, width = self.height, self.width return height, width def load_image(self, file_path): image = Image.open(file_path).convert("RGB") image = self.crop_and_resize(image, *self.get_height_width(image)) return image def load_data(self, file_path): return self.load_image(file_path) def __getitem__(self, data_id): data = self.data[data_id % len(self.data)].copy() for key in self.data_file_keys: if key in data: if isinstance(data[key], list): path = [os.path.join(self.base_path, p) for p in data[key]] data[key] = [self.load_data(p) for p in path] else: path = os.path.join(self.base_path, data[key]) data[key] = self.load_data(path) if data[key] is None: warnings.warn(f"cannot load file {data[key]}.") return None return data def __len__(self): return len(self.data) * self.repeat class VideoDataset(torch.utils.data.Dataset): def __init__( self, base_path=None, metadata_path=None, num_frames=81, time_division_factor=4, time_division_remainder=1, max_pixels=1920*1080, height=None, width=None, height_division_factor=16, width_division_factor=16, data_file_keys=("video",), image_file_extension=("jpg", "jpeg", "png", "webp"), video_file_extension=("mp4", "avi", "mov", "wmv", "mkv", "flv", "webm"), repeat=1, args=None, ): if args is not None: base_path = args.dataset_base_path metadata_path = args.dataset_metadata_path height = args.height width = args.width max_pixels = args.max_pixels num_frames = args.num_frames data_file_keys = args.data_file_keys.split(",") repeat = args.dataset_repeat self.base_path = base_path self.num_frames = num_frames self.time_division_factor = time_division_factor self.time_division_remainder = time_division_remainder self.max_pixels = max_pixels self.height = height self.width = width self.height_division_factor = height_division_factor self.width_division_factor = width_division_factor self.data_file_keys = data_file_keys self.image_file_extension = image_file_extension self.video_file_extension = video_file_extension self.repeat = repeat if height is not None and width is not None: print("Height and width are fixed. Setting `dynamic_resolution` to False.") self.dynamic_resolution = False elif height is None and width is None: print("Height and width are none. Setting `dynamic_resolution` to True.") self.dynamic_resolution = True if metadata_path is None: print("No metadata. Trying to generate it.") metadata = self.generate_metadata(base_path) print(f"{len(metadata)} lines in metadata.") self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] elif metadata_path.endswith(".json"): with open(metadata_path, "r") as f: metadata = json.load(f) self.data = metadata else: metadata = pd.read_csv(metadata_path) self.data = [metadata.iloc[i].to_dict() for i in range(len(metadata))] def generate_metadata(self, folder): video_list, prompt_list = [], [] file_set = set(os.listdir(folder)) for file_name in file_set: if "." not in file_name: continue file_ext_name = file_name.split(".")[-1].lower() file_base_name = file_name[:-len(file_ext_name)-1] if file_ext_name not in self.image_file_extension and file_ext_name not in self.video_file_extension: continue prompt_file_name = file_base_name + ".txt" if prompt_file_name not in file_set: continue with open(os.path.join(folder, prompt_file_name), "r", encoding="utf-8") as f: prompt = f.read().strip() video_list.append(file_name) prompt_list.append(prompt) metadata = pd.DataFrame() metadata["video"] = video_list metadata["prompt"] = prompt_list return metadata def crop_and_resize(self, image, target_height, target_width): width, height = image.size scale = max(target_width / width, target_height / height) image = torchvision.transforms.functional.resize( image, (round(height*scale), round(width*scale)), interpolation=torchvision.transforms.InterpolationMode.BILINEAR ) image = torchvision.transforms.functional.center_crop(image, (target_height, target_width)) return image def get_height_width(self, image): if self.dynamic_resolution: width, height = image.size if width * height > self.max_pixels: scale = (width * height / self.max_pixels) ** 0.5 height, width = int(height / scale), int(width / scale) height = height // self.height_division_factor * self.height_division_factor width = width // self.width_division_factor * self.width_division_factor else: height, width = self.height, self.width return height, width def get_num_frames(self, reader): num_frames = self.num_frames if int(reader.count_frames()) < num_frames: num_frames = int(reader.count_frames()) while num_frames > 1 and num_frames % self.time_division_factor != self.time_division_remainder: num_frames -= 1 return num_frames def load_video(self, file_path): reader = imageio.get_reader(file_path) num_frames = self.get_num_frames(reader) frames = [] for frame_id in range(num_frames): frame = reader.get_data(frame_id) frame = Image.fromarray(frame) frame = self.crop_and_resize(frame, *self.get_height_width(frame)) frames.append(frame) reader.close() return frames def load_image(self, file_path): image = Image.open(file_path).convert("RGB") image = self.crop_and_resize(image, *self.get_height_width(image)) frames = [image] return frames def is_image(self, file_path): file_ext_name = file_path.split(".")[-1] return file_ext_name.lower() in self.image_file_extension def is_video(self, file_path): file_ext_name = file_path.split(".")[-1] return file_ext_name.lower() in self.video_file_extension def load_data(self, file_path): if self.is_image(file_path): return self.load_image(file_path) elif self.is_video(file_path): return self.load_video(file_path) else: return None def __getitem__(self, data_id): data = self.data[data_id % len(self.data)].copy() for key in self.data_file_keys: if key in data: path = os.path.join(self.base_path, data[key]) data[key] = self.load_data(path) if data[key] is None: warnings.warn(f"cannot load file {data[key]}.") return None return data def __len__(self): return len(self.data) * self.repeat class AudioDataset(torch.utils.data.Dataset): def __init__( self, base_path=None, metadata_path=None, sample_rate=None, # target sample rate; if None, keep original num_samples=None, # target number of samples; if None, keep original length min_num_samples=2048, # minimum number of samples; pad to this if shorter max_num_samples=30*44100, # maximum number of samples; center-crop if longer mono=False, data_file_keys=("audio", "audio_latent"), audio_file_extension=("wav", "mp3", "flac", "ogg", "m4a", "aac", "wma", "", "mp4", "aiff", "wv"), repeat=1, drop_prompt_prob=0.1, cache_folder=None, append_duration_suffix=False, append_duration_suffix_prob=0.5, duration_precision=1, args=None, ): if args is not None: base_path = args.dataset_base_path if metadata_path is None: metadata_path = args.dataset_metadata_path # Optional arguments; use getattr to avoid hard dependency in parsers sample_rate = getattr(args, "sample_rate", sample_rate) num_samples = getattr(args, "num_audio_samples", num_samples) min_num_samples = getattr(args, "min_num_audio_samples", min_num_samples) max_num_samples = getattr(args, "max_num_audio_samples", max_num_samples) mono = getattr(args, "mono", mono) data_file_keys = getattr(args, "data_file_keys", ",".join(data_file_keys)).split(",") repeat = args.dataset_repeat drop_prompt_prob = getattr(args, "drop_prompt_prob", drop_prompt_prob) # 0.1 cache_folder = getattr(args, "cache_folder", cache_folder) append_duration_suffix = getattr(args, "append_duration_suffix", append_duration_suffix) append_duration_suffix_prob = getattr(args, "append_duration_suffix_prob", append_duration_suffix_prob) duration_precision = getattr(args, "duration_precision", duration_precision) self.base_path = base_path self.sample_rate = sample_rate self.num_samples = num_samples self.min_num_samples = min_num_samples self.max_num_samples = max_num_samples self.mono = mono self.data_file_keys = data_file_keys self.audio_file_extension = audio_file_extension self.repeat = repeat self.drop_prompt_prob = drop_prompt_prob self.cache_folder = cache_folder self.append_duration_suffix = append_duration_suffix if append_duration_suffix_prob is None: append_duration_suffix_prob = 1.0 self.append_duration_suffix_prob = float(append_duration_suffix_prob) # Clamp to [0, 1] to avoid surprises self.append_duration_suffix_prob = max(0.0, min(1.0, self.append_duration_suffix_prob)) self.duration_precision = int(duration_precision) if duration_precision is not None else 1 self._sharded_reader = None self.max_duration_s = self.max_num_samples / self.sample_rate if metadata_path is None: print("No metadata. Trying to generate it.") metadata = self.generate_metadata(base_path) print(f"{len(metadata)} lines in metadata.") self.data = metadata.to_dict("records") elif metadata_path.endswith(".json"): with open(metadata_path, "r") as f: metadata = json.load(f) self.data = metadata elif metadata_path.endswith(".jsonl"): metadata = [] with open(metadata_path, 'r') as f: for line in tqdm(f): metadata.append(json.loads(line.strip())) self.data = metadata else: metadata = pd.read_csv(metadata_path, dtype="string") self.data = metadata.to_dict("records") def generate_metadata(self, folder): audio_list, prompt_list = [], [] file_set = set(os.listdir(folder)) for file_name in file_set: if "." not in file_name: continue file_ext_name = file_name.split(".")[-1].lower() file_base_name = file_name[:-len(file_ext_name)-1] if file_ext_name not in self.audio_file_extension: continue prompt_file_name = file_base_name + ".txt" if prompt_file_name not in file_set: continue with open(os.path.join(folder, prompt_file_name), "r", encoding="utf-8") as f: prompt = f.read().strip() audio_list.append(file_name) prompt_list.append(prompt) metadata = pd.DataFrame() metadata["audio"] = audio_list metadata["prompt"] = prompt_list return metadata def is_audio(self, file_path): file_ext_name = file_path.split(".")[-1] return file_ext_name.lower() in self.audio_file_extension def _ensure_mono(self, waveform): # waveform: Tensor [C, T] or [T] if waveform.ndim == 1: waveform = waveform.unsqueeze(0) if self.mono: if waveform.shape[0] > 1: waveform = waveform.mean(dim=0, keepdim=True) else: # If mono=False but input is single-channel, duplicate to stereo if waveform.shape[0] == 1: waveform = waveform.repeat(2, 1) elif waveform.shape[0] > 2: # Some files have 6+ channels; keep only the first two. waveform = waveform[0:2] return waveform def _apply_num_samples(self, waveform): if self.num_samples is None: return waveform target = int(self.num_samples) num_current = waveform.shape[-1] if num_current == target: return waveform if num_current > target: return waveform[..., :target] # pad time dimension to target length pad_T = target - num_current pad = torch.zeros((waveform.shape[0], pad_T), dtype=waveform.dtype, device=waveform.device) return torch.cat([waveform, pad], dim=-1) def _ensure_max_num_samples(self, waveform): if self.max_num_samples is None: return waveform max_target = int(self.max_num_samples) num_current = waveform.shape[-1] if num_current <= max_target: return waveform start = (num_current - max_target) // 2 end = start + max_target return waveform[..., start:end] def _ensure_min_num_samples(self, waveform): if self.min_num_samples is None: return waveform min_target = int(self.min_num_samples) num_current = waveform.shape[-1] if num_current >= min_target: return waveform pad_T = min_target - num_current pad = torch.zeros((waveform.shape[0], pad_T), dtype=waveform.dtype, device=waveform.device) return torch.cat([waveform, pad], dim=-1) def load_audio(self, file_path, start_time=None, end_time=None): if not self.is_audio(file_path): raise ValueError(f"File {file_path} is not an audio file.") try: if not os.path.exists(file_path): warnings.warn(f"audio file not found: {file_path}.") return None # Decode with soundfile first; fall back to ffmpeg -> wav for # formats it cannot read (e.g. some mp3). Avoids torchcodec entirely # to dodge segfaults seen on certain builds. import soundfile as sf try: data, sr = sf.read(file_path, dtype='float32') waveform = torch.from_numpy(data) if waveform.ndim == 1: waveform = waveform.unsqueeze(0) else: waveform = waveform.t() waveform = waveform.float() except Exception: # soundfile cannot decode this file (e.g. some mp3 variants); # transcode to wav with ffmpeg, then re-read. import subprocess target_sr = int(self.sample_rate) if self.sample_rate is not None else 48000 target_ac = 1 if self.mono else 2 wav_tmp = file_path + ".sf_fallback.wav" subprocess.run( ["ffmpeg", "-y", "-i", file_path, "-ar", str(target_sr), "-ac", str(target_ac), "-f", "wav", "-acodec", "pcm_f32le", wav_tmp], capture_output=True, check=True, ) data, sr = sf.read(wav_tmp, dtype='float32') os.unlink(wav_tmp) waveform = torch.from_numpy(data) if waveform.ndim == 1: waveform = waveform.unsqueeze(0) else: waveform = waveform.t() waveform = waveform.float() duration_seconds = waveform.shape[-1] / sr # When start/end times are given and valid, read that segment. use_segment = (start_time is not None and end_time is not None) if use_segment: s = max(0, int(start_time * sr)) e = min(waveform.shape[-1], int(end_time * sr)) waveform = waveform[:, s:e] elif (self.max_num_samples is not None and self.sample_rate is not None and duration_seconds > self.max_duration_s): # Center-crop to max length at decode time for long samples. max_dur = float(self.max_duration_s) start_c = max(0.0, float(duration_seconds - max_dur) / 2.0) # Cap the start offset to avoid pathological decode time on # extremely long audio (observed on some freesound files). start_c = min(start_c, 60.0) end_c = start_c + max_dur s = int(start_c * sr) e = int(end_c * sr) waveform = waveform[:, s:e] elif duration_seconds > 60*60: warnings.warn(f"Duration of {file_path} is {duration_seconds} seconds, which is longer than 60 minutes.") waveform = waveform[:, :int(60*60*sr)] if self.sample_rate is not None and sr != self.sample_rate: waveform = torchaudio.functional.resample(waveform, sr, self.sample_rate) waveform = self._ensure_mono(waveform) # ensure [C, T] # Apply max-length center crop first. waveform = self._ensure_max_num_samples(waveform) # center-crop if longer than max # Compute true non-pad duration (seconds) before any padding. effective_sr = self.sample_rate if self.sample_rate is not None else sr num_current = waveform.shape[-1] if self.num_samples is None: valid_T = num_current else: target = int(self.num_samples) # If it would be cropped, take target; if it would be padded, keep original. valid_T = min(num_current, target) # Then apply exact-length and min-length pad/crop. waveform = self._apply_num_samples(waveform) # crop/pad on T to exact num_samples if provided waveform = self._ensure_min_num_samples(waveform) # ensure at least min_num_samples nonpad_duration_s = float(valid_T) / float(effective_sr) return waveform, nonpad_duration_s except Exception: warnings.warn(f"cannot load audio file {file_path}.") return None def load_data(self, file_path, start_time=None, end_time=None): return self.load_audio(file_path, start_time=start_time, end_time=end_time) def __getitem__(self, data_id): import warnings # Prefer the cache, if any. if self.cache_folder is not None: # New multi-shard bin+idx layout (enabled when the meta file exists). meta_path = os.path.join(self.cache_folder, DEFAULT_META_FILENAME) if os.path.exists(meta_path): if self._sharded_reader is None: self._sharded_reader = ShardedBinIdxReader(self.cache_folder) loaded = self._sharded_reader.get(int(data_id)) if not isinstance(loaded, dict): warnings.warn(f"cache miss for data_id={data_id}, skipping.") return None # context has shape [1, T, D]; right-pad T to 512 with zeros. context = loaded["context"] pad_T = 512 - context.shape[1] pad = torch.zeros((1, pad_T, context.shape[2]), dtype=context.dtype, device=context.device) context = torch.cat([context, pad], dim=1) loaded["context"] = context return {"cached": torch.tensor(True), "data_id": data_id, **loaded} # Legacy single-file cache layout. cache_path_npz = os.path.join(self.cache_folder, f"{data_id}.npz") cache_path_pth = os.path.join(self.cache_folder, f"{data_id}.pth") if os.path.exists(cache_path_npz) or os.path.exists(cache_path_pth): try: if os.path.exists(cache_path_npz): npz_file = np.load(cache_path_npz) loaded = {k: torch.from_numpy(np.array(v)) if hasattr(v, "dtype") else v for k, v in npz_file.items()} else: loaded = torch.load(cache_path_pth, map_location="cpu") if isinstance(loaded, dict): return {"cached": torch.tensor(True), "data_id": data_id, **loaded} except Exception: warnings.warn(f"cannot load cache file {cache_path_npz if os.path.exists(cache_path_npz) else cache_path_pth}.") pass warnings.warn(f"cannot load cache file {cache_path_npz if os.path.exists(cache_path_npz) else cache_path_pth}.") data = self.data[data_id % len(self.data)].copy() computed_duration_s = None data['data_id'] = data_id # Read start_time / end_time from the row (if present). Use to_numeric # so non-numeric strings yield NaN instead of raising. start_time_s = None end_time_s = None if "start_time" in data: val = pd.to_numeric(data.get("start_time"), errors="coerce") if pd.notna(val): start_time_s = float(val) if "end_time" in data: val = pd.to_numeric(data.get("end_time"), errors="coerce") if pd.notna(val): end_time_s = float(val) # If audio_latent is provided, load it from safetensors and build a # dummy audio tensor of matching length. latent_loaded = False latent_T = None if "audio_latent" in data and pd.notna(data["audio_latent"]): try: latent_path = os.path.join(self.base_path, data["audio_latent"]) latent_obj = safe_load_file(latent_path) if "latents" not in latent_obj: raise ValueError("key 'latents' not found in safetensors file") latents = latent_obj["latents"] if isinstance(latents, np.ndarray): latents = torch.from_numpy(latents) latents = latents.float() if latents.ndim != 2: raise ValueError(f"audio_latent expected shape [64|128, T], got {tuple(latents.shape)}") channels = int(latents.shape[0]) if channels not in (64, 128): raise ValueError(f"audio_latent expected shape [64|128, T], got {tuple(latents.shape)}") data["audio_latent"] = latents latent_T = int(latents.shape[1]) # Dummy audio: 64-ch latent -> [2, T*2048]; 128-ch -> [1, T*960]. if channels == 64: data["audio"] = torch.zeros((2, latent_T * 2048), dtype=torch.float32) else: # channels == 128 data["audio"] = torch.zeros((1, latent_T * 960), dtype=torch.float32) latent_loaded = True except Exception as e: warnings.warn(f"cannot load audio_latent file {data.get('audio_latent')}: {e}.") return None else: if "audio_latent" in data: del data["audio_latent"] for key in self.data_file_keys: if key in data: # Skip keys already populated above. if key == "audio_latent" and latent_loaded: continue if key == "audio" and latent_loaded: # audio already replaced by a dummy tensor from audio_latent. continue path = data[key] if key == "audio": loaded = self.load_data(path, start_time=start_time_s, end_time=end_time_s) if isinstance(loaded, tuple) and len(loaded) == 2: data[key], computed_duration_s = loaded else: data[key] = loaded else: loaded = self.load_data(path) data[key] = loaded[0] if isinstance(loaded, tuple) else loaded if data[key] is None: warnings.warn(f"cannot load file for key={key}.") return None # Duration is now sourced from load_audio's return value only. # Randomly drop prompt with given probability if "prompt" in data and isinstance(data["prompt"], str): if self.drop_prompt_prob is not None and self.drop_prompt_prob > 0: if random.random() < float(self.drop_prompt_prob): data["prompt"] = "" else: warnings.warn(f"prompt is not a string: {data['prompt']}.") # Optionally append a duration suffix to the prompt. if self.append_duration_suffix: # Probability of appending; default 0.5. 1.0 means always append. prob = getattr(self, "append_duration_suffix_prob", 1.0) if prob is None: prob = 1.0 prob = max(0.0, min(1.0, float(prob))) if computed_duration_s is not None: if prob >= 1.0 or (prob > 0 and random.random() < prob): computed_duration_s = min(computed_duration_s, self.max_duration_s) fmt = f"{{:.{max(0, int(self.duration_precision))}f}}" duration_text = fmt.format(computed_duration_s) suffix = f" duration: {duration_text}s" if "prompt" in data and isinstance(data["prompt"], str): data["prompt"] = data["prompt"] + suffix elif not latent_loaded and prob > 0: warnings.warn("duration info unavailable; skip duration suffix.") return data def __len__(self): return len(self.data) * self.repeat class DiffusionTrainingModule(torch.nn.Module): def __init__(self): super().__init__() def to(self, *args, **kwargs): for name, model in self.named_children(): model.to(*args, **kwargs) return self def trainable_modules(self): trainable_modules = filter(lambda p: p.requires_grad, self.parameters()) return trainable_modules def trainable_param_names(self): trainable_param_names = list(filter(lambda named_param: named_param[1].requires_grad, self.named_parameters())) trainable_param_names = set([named_param[0] for named_param in trainable_param_names]) return trainable_param_names def add_lora_to_model(self, model, target_modules, lora_rank, lora_alpha=None): if lora_alpha is None: lora_alpha = lora_rank lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha, target_modules=target_modules) model = inject_adapter_in_model(lora_config, model) return model def mapping_lora_state_dict(self, state_dict): new_state_dict = {} for key, value in state_dict.items(): if "lora_A.weight" in key or "lora_B.weight" in key: new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight") new_state_dict[new_key] = value return new_state_dict def export_trainable_state_dict(self, state_dict, remove_prefix=None): trainable_param_names = self.trainable_param_names() state_dict = {name: param for name, param in state_dict.items() if name in trainable_param_names} if remove_prefix is not None: state_dict_ = {} for name, param in state_dict.items(): if name.startswith(remove_prefix): name = name[len(remove_prefix):] state_dict_[name] = param state_dict = state_dict_ return state_dict class ModelLogger: def __init__(self, output_path, remove_prefix_in_ckpt=None, state_dict_converter=lambda x:x): self.output_path = output_path self.remove_prefix_in_ckpt = remove_prefix_in_ckpt self.state_dict_converter = state_dict_converter self.num_steps = 0 def on_step_end(self, accelerator, model, save_steps=None, loss=None): self.num_steps += 1 if save_steps is not None and self.num_steps % save_steps == 0: self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors") def on_epoch_end(self, accelerator, model, epoch_id): accelerator.wait_for_everyone() if accelerator.is_main_process: state_dict = accelerator.get_state_dict(model) state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt) state_dict = self.state_dict_converter(state_dict) os.makedirs(self.output_path, exist_ok=True) path = os.path.join(self.output_path, f"epoch-{epoch_id}.safetensors") accelerator.save(state_dict, path, safe_serialization=True) def on_training_end(self, accelerator, model, save_steps=None): if save_steps is not None and self.num_steps % save_steps != 0: self.save_model(accelerator, model, f"step-{self.num_steps}.safetensors") def save_model(self, accelerator, model, file_name): accelerator.wait_for_everyone() if accelerator.is_main_process: state_dict = accelerator.get_state_dict(model) state_dict = accelerator.unwrap_model(model).export_trainable_state_dict(state_dict, remove_prefix=self.remove_prefix_in_ckpt) state_dict = self.state_dict_converter(state_dict) os.makedirs(self.output_path, exist_ok=True) path = os.path.join(self.output_path, file_name) accelerator.save(state_dict, path, safe_serialization=True) def save_training_state(self, accelerator, epoch_id, global_step, micro_step): """Save the full training state to output_path/training_state/ for resume.""" accelerator.wait_for_everyone() state_dir = os.path.join(self.output_path, "training_state") # Save the full accelerator state (model, optimizer, scheduler, RNG). accelerator.save_state(state_dir) # Then save extra metadata (epoch, step, num_steps). if accelerator.is_main_process: metadata = { "epoch_id": epoch_id, # optimizer step; only incremented when accelerator.sync_gradients is True. "global_step": global_step, "micro_step": micro_step, # micro step; incremented on every on_step_end call. # global_step = num_steps / gradient_accumulation_steps "num_steps": self.num_steps, } metadata_path = os.path.join(state_dir, "metadata.json") with open(metadata_path, "w") as f: json.dump(metadata, f) @staticmethod def load_training_metadata(state_dir): metadata_path = os.path.join(state_dir, "metadata.json") with open(metadata_path, "r") as f: return json.load(f) def launch_training_task( dataset: torch.utils.data.Dataset, model: DiffusionTrainingModule, model_logger: ModelLogger, optimizer: torch.optim.Optimizer, scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None, batch_size: int = 1, clip_grad_norm: float = 1.0, num_workers: int = 8, save_steps: int = None, num_epochs: int = 1, gradient_accumulation_steps: int = 1, find_unused_parameters: bool = False, log_dir: Optional[str] = None, prefetch_factor: int = 2, resume_from=None, # Resume directory (points to a previous output_path). ): def collate_skip_none(batch): batch = [b for b in batch if b is not None] if len(batch) == 0: return None return torch.utils.data.dataloader.default_collate(batch) dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers, prefetch_factor=prefetch_factor, in_order=True, collate_fn=collate_skip_none ) # Enable logging with Accelerator if log_dir is provided log_with = None if log_dir is not None: log_with = "tensorboard" accelerator = Accelerator( gradient_accumulation_steps=gradient_accumulation_steps, kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=find_unused_parameters)], log_with=log_with, project_dir=log_dir, ) from accelerate.utils import set_seed set_seed(42, device_specific=True) # Build the scheduler before accelerator.prepare so world_size is known. if scheduler is None: steps_per_epoch = math.ceil(len(dataset) / max(1, batch_size)) optim_steps_per_epoch = math.ceil(steps_per_epoch / max(1, gradient_accumulation_steps)) total_steps = max(1, optim_steps_per_epoch * max(1, num_epochs)) warmup_steps = min(100, max(0, total_steps - 1)) decay_steps = max(1, int(round(total_steps * 0.10))) stable_steps = max(0, total_steps - warmup_steps - decay_steps) scheduler = get_wsd_schedule( optimizer=optimizer, num_warmup_steps=warmup_steps, num_stable_steps=stable_steps, num_decay_steps=decay_steps, ) # accelerator.prepare wraps the model in DDP (for multi-GPU), shards the # dataloader across ranks, and handles mixed-precision casting. model, optimizer, dataloader, scheduler = accelerator.prepare(model, optimizer, dataloader, scheduler) global_step = 0 micro_step = 0 # --- Resume Logic --- start_epoch = 0 if resume_from is not None: state_dir = os.path.join(resume_from, "training_state") accelerator.load_state(state_dir) metadata = ModelLogger.load_training_metadata(state_dir) # Resume from the next epoch after the one we last saved (save runs at epoch end). start_epoch = metadata["epoch_id"] + 1 global_step = metadata["global_step"] micro_step = metadata["micro_step"] model_logger.num_steps = metadata["num_steps"] if accelerator.is_main_process: print(f"Resumed from epoch {metadata['epoch_id']}, global_step {global_step}, " f"micro_step {micro_step}, continuing from epoch {start_epoch}") if accelerator.log_with is not None: accelerator.init_trackers( "training", config={ "batch_size": batch_size, "num_epochs": num_epochs, "gradient_accumulation_steps": gradient_accumulation_steps, }, ) prev_iter_end_time = time.time() prev_log_time = time.time() for epoch_id in range(start_epoch, num_epochs): for data in tqdm(dataloader, disable=not accelerator.is_main_process): if data is None: prev_iter_end_time = time.time() continue # Measure data loading time # accelerator.wait_for_everyone() data_loading_end = time.time() data_loading_time = data_loading_end - prev_iter_end_time with accelerator.accumulate(model): compute_start_time = time.time() optimizer.zero_grad() loss = model(data) accelerator.backward(loss) if accelerator.sync_gradients: grad_norm = accelerator.clip_grad_norm_(model.parameters(), clip_grad_norm) optimizer.step() model_logger.on_step_end(accelerator, model, save_steps, loss) scheduler.step() compute_end_time = time.time() compute_time = compute_end_time - compute_start_time # Log time metrics for every micro step if accelerator.log_with is not None: micro_step += 1 cur_log_time = time.time() accelerator.log( { "time/data_loading": float(data_loading_time), "time/compute": float(compute_time), "time/step": float(cur_log_time - prev_log_time), "train/epoch": epoch_id, }, step=micro_step, ) prev_log_time = cur_log_time # Log loss/lr only when optimizer actually steps if accelerator.log_with is not None and accelerator.sync_gradients: global_step += 1 # Reduce loss across processes for a global mean loss_to_log = loss.detach() reduced_loss = accelerator.reduce(loss_to_log, reduction="mean") accelerator.log( { "train/loss": float(reduced_loss.float().item()), "train/lr": float(optimizer.param_groups[0]["lr"]), "train/grad_norm": float(grad_norm), }, step=global_step, ) # if micro_step == 1000: # torch.cuda.cudart().cudaProfilerStart() # torch.autograd.profiler.emit_nvtx(record_shapes=False).__enter__() # elif micro_step == 1020: # torch.cuda.cudart().cudaProfilerStop() prev_iter_end_time = time.time() if save_steps is None: model_logger.on_epoch_end(accelerator, model, epoch_id) model_logger.save_training_state(accelerator, epoch_id, global_step, micro_step) model_logger.on_training_end(accelerator, model, save_steps) accelerator.end_training() def launch_data_process_task(model: DiffusionTrainingModule, dataset, cache_folder, num_shards: int = 16, skip_first_batches: int = 0, num_workers: int = 0, prefetch_factor: int = 4): accelerator = Accelerator() def collate_first_valid(batch): for b in batch: if b is not None: return b return {"__skip__": torch.tensor(True)} dataloader = torch.utils.data.DataLoader( dataset, # sampler=sampler, shuffle=False, # collate_fn=lambda x: x[0], collate_fn=collate_first_valid, # num_workers=8, num_workers=num_workers, pin_memory=True, # prefetch_factor=4, prefetch_factor=prefetch_factor if num_workers > 0 else None, ) model, dataloader = accelerator.prepare(model, dataloader) if skip_first_batches and skip_first_batches > 0: dataloader = accelerator.skip_first_batches(dataloader, int(skip_first_batches)) os.makedirs(cache_folder, exist_ok=True) writer = ShardedBinIdxWriter(cache_folder, num_shards=num_shards) prev_step_end_time = time.time() for i, data in enumerate(tqdm(dataloader, disable=not accelerator.is_local_main_process)): if isinstance(data, dict) and data.get("__skip__", None) is not None: accelerator.wait_for_everyone() prev_step_end_time = time.time() continue step_start_time = time.time() data_id = data['data_id'] with torch.no_grad(): unwrapped_model = model.module if hasattr(model, "module") else model preprocess_start_time = time.time() # forward_preprocess encodes audio with VAE and prompt with the # text encoder, yielding input_latents and context. inputs = unwrapped_model.forward_preprocess(data) inputs = {key: inputs[key] for key in unwrapped_model.model_input_keys if key in inputs} context = inputs["context"] # context has shape [1, T, D]; the tail is padded with zeros from # some position onward. Trim those trailing all-zero rows. # [1, T] zero_mask = (context == 0).all(dim=-1) if zero_mask.any(): # Index of the first all-zero row. t_end = int(torch.where(zero_mask[0])[0][0]) else: t_end = context.size(1) context = context[:, :t_end] inputs["context"] = context preprocess_end_time = time.time() try: # Write {input_latents, context} for this sample to disk. writer.write_sample(int(data_id), inputs, compress=False) except ValueError as e: print(f"Skipping sample {data_id} due to error: {e}") write_end_time = time.time() # print(f"write time: {write_end_time - preprocess_end_time}") accelerator.wait_for_everyone() prev_step_end_time = time.time()