# coding=utf-8 """Shared decoder utilities.""" # ===== helpers.py ===== from argparse import Namespace import collections.abc import importlib.metadata as importlib_metadata import importlib.util import json import logging import os import random import shutil import time from dataclasses import asdict, dataclass, field from itertools import repeat from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch import torch.distributed as dist from diffusers.utils.import_utils import compare_versions from easydict import EasyDict from loguru import logger from packaging.version import parse from PIL import Image from torch.distributed import ProcessGroup from torch.utils.data import default_collate from torchvision import transforms def default(value, default_val): return default_val if value is None else value def _ntuple(n): def parse(x): if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): x = tuple(x) if len(x) == 1: x = tuple(repeat(x[0], n)) return x return tuple(repeat(x, n)) return parse to_2tuple = _ntuple(2) def readable_time(src, key=None, left_steps=None): """ Convert time seconds to a readable format: DD Days, HH Hours, MM Minutes, SS Seconds """ if hasattr(src, "timers"): assert key is not None, "key must be provided when src is a Timer object" time_repr = ( f"Average: {readable_time(src.average(key))}" f" | Elapsed: {readable_time(src.elapsed_total(key))}" ) if left_steps is not None: assert isinstance(left_steps, int), "left_steps must be int" time_repr += f" | Remain: {readable_time(src.average(key) * left_steps)}" return time_repr seconds = int(src) days, seconds = divmod(seconds, 86400) hours, seconds = divmod(seconds, 3600) minutes, seconds = divmod(seconds, 60) if days > 0: return f"{days} Days, {hours} Hours, {minutes} Minutes, {seconds} Seconds" if hours > 0: return f"{hours} Hours, {minutes} Minutes, {seconds} Seconds" if minutes > 0: return f"{minutes} Minutes, {seconds} Seconds" return f"{seconds} Seconds" def print_args(title, args): print(f'------------------------ {title} ------------------------', flush=True) str_list = [] name_max_length = max([48] + [len(arg) + 3 for arg in vars(args)]) for arg in vars(args): dots = '.' * (name_max_length - len(arg)) str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg))) for arg in sorted(str_list, key=lambda x: x.lower()): print(arg, flush=True) print(f'-------------------- end of {title} ---------------------', flush=True) @dataclass class DataClassMixin: def __repr__(self): max_key_length = max([len(k) for k in asdict(self).keys()]) return "\n".join( [f"{k:>{max_key_length}}: {v}" for k, v in asdict(self).items()] ) def to_dict(self): return asdict(self) # ===== file_utils.py ===== def safe_file(path): """Create the parent directory of a file if it does not exist.""" path = Path(path) path.parent.mkdir(exist_ok=True, parents=True) return path def safe_save_file(save_path, *args, save_fn=None, **kwargs): save_path = Path(save_path) save_path.parent.mkdir(exist_ok=True, parents=True) tmp_save_path = save_path.parent / f"temp_{save_path.name}" save_to = None try: save_to = save_fn(tmp_save_path, *args, **kwargs) shutil.copyfile(tmp_save_path, save_path) save_to = save_path tmp_save_path.unlink() except Exception as e: print(f"Failed to save to {save_path}. {type(e)}: {e}") return save_to def empty_logger(): logger = logging.getLogger("rosetta_empty_logger") logger.addHandler(logging.NullHandler()) logger.setLevel(logging.CRITICAL) return logger def save_to_csv(dataframe, save_path, append=False): save_path = Path(save_path) save_path.parent.mkdir(parents=True, exist_ok=True) if append: dataframe.to_csv(save_path, index=False, mode='a', header=not save_path.exists()) else: dataframe.to_csv(save_path, index=False) def save_to_json(save_path, results, indent=4): with open(save_path, "w") as f: json.dump(results, f, indent=indent) return save_path def is_package_version(package: str, operation: str, version: str): """ Compares the current Accelerate version to a given reference with an operation. Args: package (str): The package name to check. operation (str): A string representation of an operator, such as `">"` or `"<="` version (str): A version string """ package_available = importlib.util.find_spec(package.replace("-", "_")) is not None if not package_available: return False package_version = importlib_metadata.version(package) return compare_versions(parse(package_version), operation, version) # ===== torch_utils.py ===== PRECISION_TO_TYPE = { "fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, torch.float32: torch.float32, torch.float16: torch.float16, torch.bfloat16: torch.bfloat16, } def set_manual_seed(global_seed): random.seed(global_seed) np.random.seed(global_seed) torch.manual_seed(global_seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(global_seed) def set_reproducibility(enable, global_seed=None, benchmark=None): if global_seed is not None: set_manual_seed(global_seed) if enable: os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" torch.backends.cudnn.benchmark = (not enable) if benchmark is None else benchmark torch.backends.cudnn.deterministic = enable torch.use_deterministic_algorithms(enable) def except_collate_fn(batch, except_keys=None): if except_keys is None: return default_collate(batch) ret = {} existing_except_keys = [] for expect_key in except_keys: if expect_key in batch[0]: existing_except_keys.append(expect_key) except_keys = existing_except_keys for expect_key in except_keys: ret[expect_key] = [] for batch_index in range(len(batch)): for except_key in except_keys: ret[except_key].append(batch[batch_index][except_key]) del batch[batch_index][except_key] ret.update(default_collate(batch)) return ret class Timer(object): def __init__(self, enabled=False): self.timers = {} self.enabled = enabled def __getitem__(self, key): return self.timers[key] def start(self, name, sync=False, barrier=False): if not self.enabled: return if name not in self.timers: self.timers[name] = {"total_time": 0.0, "count": 0, "sync": sync, "barrier": barrier} else: if "start_time" in self.timers[name]: raise ValueError(f"Timer '{name}' is already running.") self.timers[name].update({"sync": sync, "barrier": barrier}) if sync: torch.cuda.synchronize() if barrier: dist.barrier() self.timers[name]["start_time"] = time.time() def stop(self, name): if not self.enabled: return if name not in self.timers or "start_time" not in self.timers[name]: raise ValueError(f"Timer '{name}' was not started.") if self.timers[name]["sync"]: torch.cuda.synchronize() if self.timers[name]["barrier"]: dist.barrier() elapsed_time = time.time() - self.timers[name]["start_time"] self.timers[name]["elapsed_time"] = elapsed_time self.timers[name]["total_time"] += elapsed_time self.timers[name]["count"] += 1 del self.timers[name]["start_time"] def elapsed(self, name): if not self.enabled: return 0.0 if name not in self.timers or "elapsed_time" not in self.timers[name]: raise ValueError(f"Timer '{name}' has no recorded elapsed time.") return self.timers[name]["elapsed_time"] def elapsed_total(self, name): if not self.enabled: return 0.0 if name not in self.timers: raise ValueError(f"Timer '{name}' has no recorded times.") return self.timers[name]["total_time"] def average(self, name): if not self.enabled: return 0.0 if name not in self.timers or self.timers[name]["count"] == 0: raise ValueError(f"Timer '{name}' has no recorded times.") return self.timers[name]["total_time"] / self.timers[name]["count"] # ===== image_base.py ===== class ImageInfo: """ Class to store image information for processing and generation. """ args: EasyDict | dict def __init__( self, image_type: str = None, image_tensor: torch.Tensor = None, image_width: int = None, image_height: int = None, token_width: int = None, token_height: int = None, image_token_length: int = None, base_size: int = None, ratio_index: int = None, face_image: Optional[Image.Image] = None, ori_image_width: int = None, ori_image_height: int = None, ): if self.args is None: raise ValueError("ImageInfo requires `args` attribute to be set.") self.image_type = image_type self.image_tensor = image_tensor self.ori_image_width = ori_image_width self.image_width = image_width self.w = image_width self.ori_image_height = ori_image_height self.image_height = image_height self.h = image_height self.token_width = token_width self.tk_w = token_width self.token_height = token_height self.tk_h = token_height self.image_token_length = default( image_token_length, token_width * token_height if (token_width is not None and token_height is not None) else None ) self.base_size = base_size self.ratio_index = ratio_index self.face_image = face_image # args self.add_timestep_token = self.args.get("add_timestep_token", False) self.add_image_shape_token = self.args.get("add_image_shape_token", False) def __getitem__(self, key: str) -> Any: """Allow dictionary-like access to attributes.""" if hasattr(self, key): return getattr(self, key) raise KeyError(f"Key '{key}' not found in ImageInfo") def __setitem__(self, key: str, value: Any) -> None: """Allow dictionary-like assignment to attributes.""" if hasattr(self, key): setattr(self, key, value) else: raise KeyError(f"Key '{key}' not found in ImageInfo") def __contains__(self, key: str) -> bool: """Check if the key exists in the ImageInfo object.""" return hasattr(self, key) def __repr__(self): return (f"ImageInfo(image_type={self.image_type}, image_tensor={self.image_tensor}, " f"ori_image_width={self.ori_image_width}, ori_image_height={self.ori_image_height}, " f"image_width={self.image_width}, image_height={self.image_height}, " f"token_width={self.token_width}, token_height={self.token_height}, " f"image_token_length={self.image_token_length}, " f"base_size={self.base_size}, ratio_index={self.ratio_index}, face_image={self.face_image})") @property def meta_info(self): if self.args is None: raise ValueError("meta_info requires `args` attribute to be set.") if self.image_type in ["vae", "src_image", "gen_image"]: return dict( token_length=self.image_token_length, add_timestep_token=self.add_timestep_token, add_image_shape_token=self.add_image_shape_token, base_size=self.base_size, ratio_idx=self.ratio_index, token_height=self.token_height, token_width=self.token_width, image_height=self.image_height, image_width=self.image_width, ori_image_width=self.ori_image_width, ori_image_height=self.ori_image_height, ) elif self.image_type in ["vision_encoder", "und_image", "qwen3vl"]: return dict( token_length=self.image_token_length, add_image_shape_token=self.add_image_shape_token, token_height=self.token_height, token_width=self.token_width, image_height=self.image_height, image_width=self.image_width, ori_image_width=self.ori_image_width, ori_image_height=self.ori_image_height, ) elif self.image_type == "face": return dict( token_length=self.image_token_length, ) else: raise ValueError(f"Unknown image type '{self.image_type}'") def copy(self, copy_image_tensor=True): if copy_image_tensor and self.image_tensor is None: raise ValueError("image_tensor is None, cannot copy") return ImageInfo( image_type=self.image_type, image_tensor=self.image_tensor.clone() if copy_image_tensor else None, image_width=self.image_width, image_height=self.image_height, ori_image_width=self.ori_image_width, ori_image_height=self.ori_image_height, token_width=self.token_width, token_height=self.token_height, image_token_length=self.image_token_length, base_size=self.base_size, ratio_index=self.ratio_index, face_image=self.face_image, # shared ) class ImageTensor(torch.Tensor): # This class is just for type hinting purposes. Attribute `i` should be defined # as an instance attribute of the torch.Tensor instance, like: tensor.i = ImageInfo(...) i: ImageInfo section_type: str vision_encoder_kwargs: dict[str, torch.Tensor] class JointImageInfo(object): def __init__(self, vae_image_info: ImageInfo, vision_image_info: ImageInfo, vision_encoder_kwargs: dict = None): self.vae_image_info = vae_image_info self.vision_image_info = vision_image_info self.vision_encoder_kwargs = vision_encoder_kwargs # Define key attributes to align with ImageInfo for uniformity self.image_type = "joint_image" self.image_token_length = vae_image_info.image_token_length + vision_image_info.image_token_length self.add_timestep_token = vae_image_info.add_timestep_token self.add_image_shape_token = vae_image_info.add_image_shape_token def __repr__(self): return f"JointImageInfo(vae_image={self.vae_image_info}, vision_image={self.vision_image_info})" @property def meta_info(self): # Used for image sections of tkwrapper.encode_general() return dict( token_length=[self.vae_image_info.image_token_length, self.vision_image_info.image_token_length], add_timestep_token=self.add_timestep_token, add_image_shape_token=self.add_image_shape_token, base_size=self.vae_image_info.base_size, ratio_idx=self.vae_image_info.ratio_index, token_height=[self.vae_image_info.token_height, self.vision_image_info.token_height], token_width=[self.vae_image_info.token_width, self.vision_image_info.token_width], image_height=[self.vae_image_info.image_height, self.vision_image_info.image_height], image_width=[self.vae_image_info.image_width, self.vision_image_info.image_width], ) def copy(self, copy_image_tensor=True): if copy_image_tensor and (self.vae_image_info.image_tensor is None or self.vision_image_info.image_tensor is None): raise ValueError("image_tensor is None, cannot copy") return JointImageInfo(self.vae_image_info.copy(copy_image_tensor), self.vision_image_info.copy(copy_image_tensor), self.vision_encoder_kwargs) class CondImage(object): def __init__(self, image_type: str, vae_image: ImageTensor, vit_image: ImageTensor): self.image_type = image_type self.vae_image = vae_image self.vit_image = vit_image if image_type == "vae": self.i = vae_image.i self.section_type = "cond_vae_image" elif image_type == "vit": self.i = vit_image.i self.section_type = "cond_vit_image" elif image_type == "vae_vit": self.i = JointImageInfo(vae_image.i, vit_image.i) self.section_type = "cond_joint_image" else: raise ValueError(f"Unknown image_type: {image_type}") # Runtime state helpers """Runtime helpers shared by training and inference.""" # Global runtime state # args _GLOBAL_ARGS: Optional[Namespace] = None # parallel state _GLOBAL_PARALLEL_STATE = None # logger _GLOBAL_LOGGER = None # Global Switch _ensure_initialized: bool = True _ensure_not_initialized: bool = True def _ensure_var_is_initialized(var, name): if not _ensure_initialized: return assert var is not None, '{} is not initialized.'.format(name) def _ensure_var_is_not_initialized(var, name): if not _ensure_not_initialized: return assert var is None, '{} is already initialized.'.format(name) def set_args(args): global _GLOBAL_ARGS _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args') _GLOBAL_ARGS = args def get_args(): global _GLOBAL_ARGS _ensure_var_is_initialized(_GLOBAL_ARGS, 'args') return _GLOBAL_ARGS def set_parallel_state(parallel_state): global _GLOBAL_PARALLEL_STATE _ensure_var_is_not_initialized(_GLOBAL_PARALLEL_STATE, "parallel_state") _GLOBAL_PARALLEL_STATE = parallel_state def get_parallel_state(): global _GLOBAL_PARALLEL_STATE _ensure_var_is_initialized(_GLOBAL_PARALLEL_STATE, "parallel_state") return _GLOBAL_PARALLEL_STATE def set_logger(logger): global _GLOBAL_LOGGER _ensure_var_is_not_initialized(_GLOBAL_LOGGER, 'logger') _GLOBAL_LOGGER = logger def get_logger(): global _GLOBAL_LOGGER _ensure_var_is_initialized(_GLOBAL_LOGGER, 'logger') return _GLOBAL_LOGGER # Parallel state @dataclass class ParallelState: # Data parallel dp_rank: int dp_size: int dp_group: ProcessGroup = None # Tensor parallel tp_rank: int = 0 tp_size: int = 1 # Pipeline parallel pp_rank: int = 0 pp_size: int = 1 # Expert parallel ep_rank: int = 0 ep_size: int = 1 # Context parallel cp_rank: int = 0 cp_size: int = 1 def __post_init__(self): set_parallel_state(self) def __repr__(self): res = [] res.append(f"dp: {self.dp_rank}/{self.dp_size}") res.append(f"tp: {self.tp_rank}/{self.tp_size}") res.append(f"pp: {self.pp_rank}/{self.pp_size}") res.append(f"cp: {self.cp_rank}/{self.cp_size}") return ", ".join(res) @staticmethod def from_pure_torch(): """Initialize ParallelState from native torch.distributed. For inference with pure data parallelism: each rank holds the full model, dp_rank = global rank, dp_size = world_size. """ rank = dist.get_rank() world_size = dist.get_world_size() return ParallelState( dp_rank=rank, dp_size=world_size, dp_group=None, ) def get_tensor_and_data_parallel_group(self, with_context_parallel: bool = False): return dist.GroupMember.WORLD