feat: Implement self-contained specialist managers

common/README.md

@@ -0,0 +1,138 @@
+# 🛠️ helpers/ - Ferramentas de IA de Terceiros Adaptadas para ADUC-SDR
+
+Esta pasta contém implementações adaptadas de modelos e utilitários de IA de terceiros, que servem como "especialistas" ou "ferramentas" de baixo nível para a arquitetura ADUC-SDR.
+
+**IMPORTANTE:** O conteúdo desta pasta é de autoria de seus respectivos idealizadores e desenvolvedores originais. Esta pasta **NÃO FAZ PARTE** do projeto principal ADUC-SDR em termos de sua arquitetura inovadora. Ela serve como um repositório para as **dependências diretas e modificadas** que os `DeformesXDEngines` (os estágios do "foguete" ADUC-SDR) invocam para realizar tarefas específicas (geração de imagem, vídeo, áudio).
+
+As modificações realizadas nos arquivos aqui presentes visam principalmente:
+1.  **Adaptação de Interfaces:** Padronizar as interfaces para que se encaixem no fluxo de orquestração do ADUC-SDR.
+2.  **Gerenciamento de Recursos:** Integrar lógicas de carregamento/descarregamento de modelos (GPU management) e configurações via arquivos YAML.
+3.  **Otimização de Fluxo:** Ajustar as pipelines para aceitar formatos de entrada mais eficientes (ex: tensores pré-codificados em vez de caminhos de mídia, pulando etapas de codificação/decodificação redundantes).
+
+---
+
+## 📄 Licenciamento
+
+O conteúdo original dos projetos listados abaixo é licenciado sob a **Licença Apache 2.0**, ou outra licença especificada pelos autores originais. Todas as modificações e o uso desses arquivos dentro da estrutura `helpers/` do projeto ADUC-SDR estão em conformidade com os termos da **Licença Apache 2.0**.
+
+As licenças originais dos projetos podem ser encontradas nas suas respectivas fontes ou nos subdiretórios `incl_licenses/` dentro de cada módulo adaptado.
+
+---
+
+## 🛠️ API dos Helpers e Guia de Uso
+
+Esta seção detalha como cada helper (agente especialista) deve ser utilizado dentro do ecossistema ADUC-SDR. Todos os agentes são instanciados como **singletons** no `hardware_manager.py` para garantir o gerenciamento centralizado de recursos de GPU.
+
+### **gemini_helpers.py (GeminiAgent)**
+
+*   **Propósito:** Atua como o "Oráculo de Síntese Adaptativo", responsável por todas as tarefas de processamento de linguagem natural, como criação de storyboards, geração de prompts, e tomada de decisões narrativas.
+*   **Singleton Instance:** `gemini_agent_singleton`
+*   **Construtor:** `GeminiAgent()`
+    *   Lê `configs/gemini_config.yaml` para obter o nome do modelo, parâmetros de inferência e caminhos de templates de prompt. A chave da API é lida da variável de ambiente `GEMINI_API_KEY`.
+*   **Métodos Públicos:**
+    *   `generate_storyboard(prompt: str, num_keyframes: int, ref_image_paths: list[str])`
+        *   **Inputs:**
+            *   `prompt`: A ideia geral do filme (string).
+            *   `num_keyframes`: O número de cenas a serem geradas (int).
+            *   `ref_image_paths`: Lista de caminhos para as imagens de referência (list[str]).
+        *   **Output:** `tuple[list[str], str]` (Uma tupla contendo a lista de strings do storyboard e um relatório textual da operação).
+    *   `select_keyframes_from_pool(storyboard: list, base_image_paths: list[str], pool_image_paths: list[str])`
+        *   **Inputs:**
+            *   `storyboard`: A lista de strings do storyboard gerado.
+            *   `base_image_paths`: Imagens de referência base (list[str]).
+            *   `pool_image_paths`: O "banco de imagens" de onde selecionar (list[str]).
+        *   **Output:** `tuple[list[str], str]` (Uma tupla contendo a lista de caminhos de imagens selecionadas e um relatório textual).
+    *   `get_anticipatory_keyframe_prompt(...)`
+        *   **Inputs:** Contexto narrativo e visual para gerar um prompt de imagem.
+        *   **Output:** `tuple[str, str]` (Uma tupla contendo o prompt gerado para o modelo de imagem e um relatório textual).
+    *   `get_initial_motion_prompt(...)`
+        *   **Inputs:** Contexto narrativo e visual para a primeira transição de vídeo.
+        *   **Output:** `tuple[str, str]` (Uma tupla contendo o prompt de movimento gerado e um relatório textual).
+    *   `get_transition_decision(...)`
+        *   **Inputs:** Contexto narrativo e visual para uma transição de vídeo intermediária.
+        *   **Output:** `tuple[dict, str]` (Uma tupla contendo um dicionário `{"transition_type": "...", "motion_prompt": "..."}` e um relatório textual).
+    *   `generate_audio_prompts(...)`
+        *   **Inputs:** Contexto narrativo global.
+        *   **Output:** `tuple[dict, str]` (Uma tupla contendo um dicionário `{"music_prompt": "...", "sfx_prompt": "..."}` e um relatório textual).
+
+### **flux_kontext_helpers.py (FluxPoolManager)**
+
+*   **Propósito:** Especialista em geração de imagens de alta qualidade (keyframes) usando a pipeline FluxKontext. Gerencia um pool de workers para otimizar o uso de múltiplas GPUs.
+*   **Singleton Instance:** `flux_kontext_singleton`
+*   **Construtor:** `FluxPoolManager(device_ids: list[str], flux_config_file: str)`
+    *   Lê `configs/flux_config.yaml`.
+*   **Método Público:**
+    *   `generate_image(prompt: str, reference_images: list[Image.Image], width: int, height: int, seed: int = 42, callback: callable = None)`
+        *   **Inputs:**
+            *   `prompt`: Prompt textual para guiar a geração (string).
+            *   `reference_images`: Lista de objetos `PIL.Image` como referência visual.
+            *   `width`, `height`: Dimensões da imagem de saída (int).
+            *   `seed`: Semente para reprodutibilidade (int).
+            *   `callback`: Função de callback opcional para monitorar o progresso.
+        *   **Output:** `PIL.Image.Image` (O objeto da imagem gerada).
+
+### **dreamo_helpers.py (DreamOAgent)**
+
+*   **Propósito:** Especialista em geração de imagens de alta qualidade (keyframes) usando a pipeline DreamO, com capacidades avançadas de edição e estilo a partir de referências.
+*   **Singleton Instance:** `dreamo_agent_singleton`
+*   **Construtor:** `DreamOAgent(device_id: str = None)`
+    *   Lê `configs/dreamo_config.yaml`.
+*   **Método Público:**
+    *   `generate_image(prompt: str, reference_images: list[Image.Image], width: int, height: int)`
+        *   **Inputs:**
+            *   `prompt`: Prompt textual para guiar a geração (string).
+            *   `reference_images`: Lista de objetos `PIL.Image` como referência visual. A lógica interna atribui a primeira imagem como `style` e as demais como `ip`.
+            *   `width`, `height`: Dimensões da imagem de saída (int).
+        *   **Output:** `PIL.Image.Image` (O objeto da imagem gerada).
+
+### **ltx_manager_helpers.py (LtxPoolManager)**
+
+*   **Propósito:** Especialista na geração de fragmentos de vídeo no espaço latente usando a pipeline LTX-Video. Gerencia um pool de workers para otimizar o uso de múltiplas GPUs.
+*   **Singleton Instance:** `ltx_manager_singleton`
+*   **Construtor:** `LtxPoolManager(device_ids: list[str], ltx_model_config_file: str, ltx_global_config_file: str)`
+    *   Lê o `ltx_global_config_file` e o `ltx_model_config_file` para configurar a pipeline.
+*   **Método Público:**
+    *   `generate_latent_fragment(**kwargs)`
+        *   **Inputs:** Dicionário de keyword arguments (`kwargs`) contendo todos os parâmetros da pipeline LTX, incluindo:
+            *   `height`, `width`: Dimensões do vídeo (int).
+            *   `video_total_frames`: Número total de frames a serem gerados (int).
+            *   `video_fps`: Frames por segundo (int).
+            *   `motion_prompt`: Prompt de movimento (string).
+            *   `conditioning_items_data`: Lista de objetos `LatentConditioningItem` contendo os tensores latentes de condição.
+            *   `guidance_scale`, `stg_scale`, `num_inference_steps`, etc.
+        *   **Output:** `tuple[torch.Tensor, tuple]` (Uma tupla contendo o tensor latente gerado e os valores de padding utilizados).
+
+### **mmaudio_helper.py (MMAudioAgent)**
+
+*   **Propósito:** Especialista em geração de áudio para um determinado fragmento de vídeo.
+*   **Singleton Instance:** `mmaudio_agent_singleton`
+*   **Construtor:** `MMAudioAgent(workspace_dir: str, device_id: str = None, mmaudio_config_file: str)`
+    *   Lê `configs/mmaudio_config.yaml`.
+*   **Método Público:**
+    *   `generate_audio_for_video(video_path: str, prompt: str, negative_prompt: str, duration_seconds: float)`
+        *   **Inputs:**
+            *   `video_path`: Caminho para o arquivo de vídeo silencioso (string).
+            *   `prompt`: Prompt textual para guiar a geração de áudio (string).
+            *   `negative_prompt`: Prompt negativo para áudio (string).
+            *   `duration_seconds`: Duração exata do vídeo (float).
+        *   **Output:** `str` (O caminho para o novo arquivo de vídeo com a faixa de áudio integrada).
+
+
+### https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/tree/main
+
+---
+
+## 🔗 Projetos Originais e Atribuições
+(A seção de atribuições e licenças permanece a mesma que definimos anteriormente)
+
+### DreamO
+*   **Repositório Original:** [https://github.com/bytedance/DreamO](https://github.com/bytedance/DreamO)
+...
+
+### LTX-Video
+*   **Repositório Original:** [https://github.com/Lightricks/LTX-Video](https://github.com/Lightricks/LTX-Video)
+...
+
+### MMAudio
+*   **Repositório Original:** [https://github.com/hkchengrex/MMAudio](https://github.com/hkchengrex/MMAudio)
+...

common/cache.py

@@ -0,0 +1,47 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+from typing import Callable
+
+
+class Cache:
+    """Caching reusable args for faster inference"""
+
+    def __init__(self, disable=False, prefix="", cache=None):
+        self.cache = cache if cache is not None else {}
+        self.disable = disable
+        self.prefix = prefix
+
+    def __call__(self, key: str, fn: Callable):
+        if self.disable:
+            return fn()
+
+        key = self.prefix + key
+        try:
+            result = self.cache[key]
+        except KeyError:
+            result = fn()
+            self.cache[key] = result
+        return result
+
+    def namespace(self, namespace: str):
+        return Cache(
+            disable=self.disable,
+            prefix=self.prefix + namespace + ".",
+            cache=self.cache,
+        )
+
+    def get(self, key: str):
+        key = self.prefix + key
+        return self.cache[key]

common/config.py

@@ -0,0 +1,110 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Configuration utility functions
+"""
+
+import importlib
+from typing import Any, Callable, List, Union
+from omegaconf import DictConfig, ListConfig, OmegaConf
+
+OmegaConf.register_new_resolver("eval", eval)
+
+
+def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]:
+    """
+    Load a configuration. Will resolve inheritance.
+    """
+    config = OmegaConf.load(path)
+    if argv is not None:
+        config_argv = OmegaConf.from_dotlist(argv)
+        config = OmegaConf.merge(config, config_argv)
+    config = resolve_recursive(config, resolve_inheritance)
+    return config
+
+
+def resolve_recursive(
+    config: Any,
+    resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]],
+) -> Any:
+    config = resolver(config)
+    if isinstance(config, DictConfig):
+        for k in config.keys():
+            v = config.get(k)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[k] = resolve_recursive(v, resolver)
+    if isinstance(config, ListConfig):
+        for i in range(len(config)):
+            v = config.get(i)
+            if isinstance(v, (DictConfig, ListConfig)):
+                config[i] = resolve_recursive(v, resolver)
+    return config
+
+
+def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any:
+    """
+    Recursively resolve inheritance if the config contains:
+    __inherit__: path/to/parent.yaml or a ListConfig of such paths.
+    """
+    if isinstance(config, DictConfig):
+        inherit = config.pop("__inherit__", None)
+
+        if inherit:
+            inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit]
+
+            parent_config = None
+            for parent_path in inherit_list:
+                assert isinstance(parent_path, str)
+                parent_config = (
+                    load_config(parent_path)
+                    if parent_config is None
+                    else OmegaConf.merge(parent_config, load_config(parent_path))
+                )
+
+            if len(config.keys()) > 0:
+                config = OmegaConf.merge(parent_config, config)
+            else:
+                config = parent_config
+    return config
+
+
+def import_item(path: str, name: str) -> Any:
+    """
+    Import a python item. Example: import_item("path.to.file", "MyClass") -> MyClass
+    """
+    return getattr(importlib.import_module(path), name)
+
+
+def create_object(config: DictConfig) -> Any:
+    """
+    Create an object from config.
+    The config is expected to contains the following:
+    __object__:
+      path: path.to.module
+      name: MyClass
+      args: as_config | as_params (default to as_config)
+    """
+    item = import_item(
+        path=config.__object__.path,
+        name=config.__object__.name,
+    )
+    args = config.__object__.get("args", "as_config")
+    if args == "as_config":
+        return item(config)
+    if args == "as_params":
+        config = OmegaConf.to_object(config)
+        config.pop("__object__")
+        return item(**config)
+    raise NotImplementedError(f"Unknown args type: {args}")

common/decorators.py

@@ -0,0 +1,147 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Decorators.
+"""
+
+import functools
+import threading
+import time
+from typing import Callable
+import torch
+
+from common.distributed import barrier_if_distributed, get_global_rank, get_local_rank
+from common.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def log_on_entry(func: Callable) -> Callable:
+    """
+    Functions with this decorator will log the function name at entry.
+    When using multiple decorators, this must be applied innermost to properly capture the name.
+    """
+
+    def log_on_entry_wrapper(*args, **kwargs):
+        logger.info(f"Entering {func.__name__}")
+        return func(*args, **kwargs)
+
+    return log_on_entry_wrapper
+
+
+def barrier_on_entry(func: Callable) -> Callable:
+    """
+    Functions with this decorator will start executing when all ranks are ready to enter.
+    """
+
+    def barrier_on_entry_wrapper(*args, **kwargs):
+        barrier_if_distributed()
+        return func(*args, **kwargs)
+
+    return barrier_on_entry_wrapper
+
+
+def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable:
+    """
+    Helper function for local_rank_zero_only and global_rank_zero_only.
+    """
+
+    def conditional_execute_wrapper(*args, **kwargs):
+        # Only execute if needed.
+        result = func(*args, **kwargs) if execute else None
+        # All GPUs must wait.
+        barrier_if_distributed()
+        # Return results.
+        return result
+
+    return conditional_execute_wrapper
+
+
+def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable:
+    """
+    Helper function for some functions with special constraints,
+    especially functions called by other global_rank_zero_only / local_rank_zero_only ones,
+    in case they are wrongly invoked in other scenarios.
+    """
+
+    def asserted_execute_wrapper(*args, **kwargs):
+        assert condition, err_msg
+        result = func(*args, **kwargs)
+        return result
+
+    return asserted_execute_wrapper
+
+
+def local_rank_zero_only(func: Callable) -> Callable:
+    """
+    Functions with this decorator will only execute on local rank zero.
+    """
+    return _conditional_execute_wrapper_factory(get_local_rank() == 0, func)
+
+
+def global_rank_zero_only(func: Callable) -> Callable:
+    """
+    Functions with this decorator will only execute on global rank zero.
+    """
+    return _conditional_execute_wrapper_factory(get_global_rank() == 0, func)
+
+
+def assert_only_global_rank_zero(func: Callable) -> Callable:
+    """
+    Functions with this decorator are only accessible to processes with global rank zero.
+    """
+    return _asserted_wrapper_factory(
+        get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0"
+    )
+
+
+def assert_only_local_rank_zero(func: Callable) -> Callable:
+    """
+    Functions with this decorator are only accessible to processes with local rank zero.
+    """
+    return _asserted_wrapper_factory(
+        get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0"
+    )
+
+
+def new_thread(func: Callable) -> Callable:
+    """
+    Functions with this decorator will run in a new thread.
+    The function will return the thread, which can be joined to wait for completion.
+    """
+
+    def new_thread_wrapper(*args, **kwargs):
+        thread = threading.Thread(target=func, args=args, kwargs=kwargs)
+        thread.start()
+        return thread
+
+    return new_thread_wrapper
+
+
+def log_runtime(func: Callable) -> Callable:
+    """
+    Functions with this decorator will logging the runtime.
+    """
+
+    @functools.wraps(func)
+    def wrapped(*args, **kwargs):
+        torch.distributed.barrier()
+        start = time.perf_counter()
+        result = func(*args, **kwargs)
+        torch.distributed.barrier()
+        logger.info(f"Completed {func.__name__} in {time.perf_counter() - start:.3f} seconds.")
+        return result
+
+    return wrapped

common/diffusion/__init__.py

@@ -0,0 +1,56 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Diffusion package.
+"""
+
+from .config import (
+    create_sampler_from_config,
+    create_sampling_timesteps_from_config,
+    create_schedule_from_config,
+)
+from .samplers.base import Sampler
+from .samplers.euler import EulerSampler
+from .schedules.base import Schedule
+from .schedules.lerp import LinearInterpolationSchedule
+from .timesteps.base import SamplingTimesteps, Timesteps
+from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
+from .types import PredictionType, SamplingDirection
+from .utils import classifier_free_guidance, classifier_free_guidance_dispatcher, expand_dims
+
+__all__ = [
+    # Configs
+    "create_sampler_from_config",
+    "create_sampling_timesteps_from_config",
+    "create_schedule_from_config",
+    # Schedules
+    "Schedule",
+    "DiscreteVariancePreservingSchedule",
+    "LinearInterpolationSchedule",
+    # Samplers
+    "Sampler",
+    "EulerSampler",
+    # Timesteps
+    "Timesteps",
+    "SamplingTimesteps",
+    # Types
+    "PredictionType",
+    "SamplingDirection",
+    "UniformTrailingSamplingTimesteps",
+    # Utils
+    "classifier_free_guidance",
+    "classifier_free_guidance_dispatcher",
+    "expand_dims",
+]

common/diffusion/config.py

@@ -0,0 +1,74 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Utility functions for creating schedules and samplers from config.
+"""
+
+import torch
+from omegaconf import DictConfig
+
+from .samplers.base import Sampler
+from .samplers.euler import EulerSampler
+from .schedules.base import Schedule
+from .schedules.lerp import LinearInterpolationSchedule
+from .timesteps.base import SamplingTimesteps
+from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps
+
+
+def create_schedule_from_config(
+    config: DictConfig,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> Schedule:
+    """
+    Create a schedule from configuration.
+    """
+    if config.type == "lerp":
+        return LinearInterpolationSchedule(T=config.get("T", 1.0))
+
+    raise NotImplementedError
+
+
+def create_sampler_from_config(
+    config: DictConfig,
+    schedule: Schedule,
+    timesteps: SamplingTimesteps,
+) -> Sampler:
+    """
+    Create a sampler from configuration.
+    """
+    if config.type == "euler":
+        return EulerSampler(
+            schedule=schedule,
+            timesteps=timesteps,
+            prediction_type=config.prediction_type,
+        )
+    raise NotImplementedError
+
+
+def create_sampling_timesteps_from_config(
+    config: DictConfig,
+    schedule: Schedule,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> SamplingTimesteps:
+    if config.type == "uniform_trailing":
+        return UniformTrailingSamplingTimesteps(
+            T=schedule.T,
+            steps=config.steps,
+            shift=config.get("shift", 1.0),
+            device=device,
+        )
+    raise NotImplementedError

common/diffusion/samplers/base.py

@@ -0,0 +1,108 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Sampler base class.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable
+import torch
+from tqdm import tqdm
+
+from ..schedules.base import Schedule
+from ..timesteps.base import SamplingTimesteps
+from ..types import PredictionType, SamplingDirection
+from ..utils import assert_schedule_timesteps_compatible
+
+
+@dataclass
+class SamplerModelArgs:
+    x_t: torch.Tensor
+    t: torch.Tensor
+    i: int
+
+
+class Sampler(ABC):
+    """
+    Samplers are ODE/SDE solvers.
+    """
+
+    def __init__(
+        self,
+        schedule: Schedule,
+        timesteps: SamplingTimesteps,
+        prediction_type: PredictionType,
+        return_endpoint: bool = True,
+    ):
+        assert_schedule_timesteps_compatible(
+            schedule=schedule,
+            timesteps=timesteps,
+        )
+        self.schedule = schedule
+        self.timesteps = timesteps
+        self.prediction_type = prediction_type
+        self.return_endpoint = return_endpoint
+
+    @abstractmethod
+    def sample(
+        self,
+        x: torch.Tensor,
+        f: Callable[[SamplerModelArgs], torch.Tensor],
+    ) -> torch.Tensor:
+        """
+        Generate a new sample given the the intial sample x and score function f.
+        """
+
+    def get_next_timestep(
+        self,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Get the next sample timestep.
+        Support multiple different timesteps t in a batch.
+        If no more steps, return out of bound value -1 or T+1.
+        """
+        T = self.timesteps.T
+        steps = len(self.timesteps)
+        curr_idx = self.timesteps.index(t)
+        next_idx = curr_idx + 1
+        bound = -1 if self.timesteps.direction == SamplingDirection.backward else T + 1
+
+        s = self.timesteps[next_idx.clamp_max(steps - 1)]
+        s = s.where(next_idx < steps, bound)
+        return s
+
+    def get_endpoint(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Get to the endpoint of the probability flow.
+        """
+        x_0, x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
+        return x_0 if self.timesteps.direction == SamplingDirection.backward else x_T
+
+    def get_progress_bar(self):
+        """
+        Get progress bar for sampling.
+        """
+        return tqdm(
+            iterable=range(len(self.timesteps) - (0 if self.return_endpoint else 1)),
+            dynamic_ncols=True,
+            desc=self.__class__.__name__,
+        )

common/diffusion/samplers/euler.py

@@ -0,0 +1,89 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+
+"""
+Euler ODE solver.
+"""
+
+from typing import Callable
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+from models.dit_v2 import na
+
+from ..types import PredictionType
+from ..utils import expand_dims
+from .base import Sampler, SamplerModelArgs
+
+
+class EulerSampler(Sampler):
+    """
+    The Euler method is the simplest ODE solver.
+    <https://en.wikipedia.org/wiki/Euler_method>
+    """
+
+    def sample(
+        self,
+        x: torch.Tensor,
+        f: Callable[[SamplerModelArgs], torch.Tensor],
+    ) -> torch.Tensor:
+        timesteps = self.timesteps.timesteps
+        progress = self.get_progress_bar()
+        i = 0
+        for t, s in zip(timesteps[:-1], timesteps[1:]):
+            pred = f(SamplerModelArgs(x, t, i))
+            x = self.step_to(pred, x, t, s)
+            i += 1
+            progress.update()
+
+        if self.return_endpoint:
+            t = timesteps[-1]
+            pred = f(SamplerModelArgs(x, t, i))
+            x = self.get_endpoint(pred, x, t)
+            progress.update()
+        return x
+
+    def step(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Step to the next timestep.
+        """
+        return self.step_to(pred, x_t, t, self.get_next_timestep(t))
+
+    def step_to(
+        self,
+        pred: torch.Tensor,
+        x_t: torch.Tensor,
+        t: torch.Tensor,
+        s: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Steps from x_t at timestep t to x_s at timestep s. Returns x_s.
+        """
+        t = expand_dims(t, x_t.ndim)
+        s = expand_dims(s, x_t.ndim)
+        T = self.schedule.T
+        # Step from x_t to x_s.
+        pred_x_0, pred_x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t)
+        pred_x_s = self.schedule.forward(pred_x_0, pred_x_T, s.clamp(0, T))
+        # Clamp x_s to x_0 and x_T if s is out of bound.
+        pred_x_s = pred_x_s.where(s >= 0, pred_x_0)
+        pred_x_s = pred_x_s.where(s <= T, pred_x_T)
+        return pred_x_s

common/diffusion/schedules/base.py

@@ -0,0 +1,131 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Schedule base class.
+"""
+
+from abc import ABC, abstractmethod, abstractproperty
+from typing import Tuple, Union
+import torch
+
+from ..types import PredictionType
+from ..utils import expand_dims
+
+
+class Schedule(ABC):
+    """
+    Diffusion schedules are uniquely defined by T, A, B:
+
+        x_t = A(t) * x_0 + B(t) * x_T, where t in [0, T]
+
+    Schedules can be continuous or discrete.
+    """
+
+    @abstractproperty
+    def T(self) -> Union[int, float]:
+        """
+        Maximum timestep inclusive.
+        Schedule is continuous if float, discrete if int.
+        """
+
+    @abstractmethod
+    def A(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Interpolation coefficient A.
+        Returns tensor with the same shape as t.
+        """
+
+    @abstractmethod
+    def B(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Interpolation coefficient B.
+        Returns tensor with the same shape as t.
+        """
+
+    # ----------------------------------------------------
+
+    def snr(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Signal to noise ratio.
+        Returns tensor with the same shape as t.
+        """
+        return (self.A(t) ** 2) / (self.B(t) ** 2)
+
+    def isnr(self, snr: torch.Tensor) -> torch.Tensor:
+        """
+        Inverse signal to noise ratio.
+        Returns tensor with the same shape as snr.
+        Subclass may implement.
+        """
+        raise NotImplementedError
+
+    # ----------------------------------------------------
+
+    def is_continuous(self) -> bool:
+        """
+        Whether the schedule is continuous.
+        """
+        return isinstance(self.T, float)
+
+    def forward(self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """
+        Diffusion forward function.
+        """
+        t = expand_dims(t, x_0.ndim)
+        return self.A(t) * x_0 + self.B(t) * x_T
+
+    def convert_from_pred(
+        self, pred: torch.Tensor, pred_type: PredictionType, x_t: torch.Tensor, t: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Convert from prediction. Return predicted x_0 and x_T.
+        """
+        t = expand_dims(t, x_t.ndim)
+        A_t = self.A(t)
+        B_t = self.B(t)
+
+        if pred_type == PredictionType.x_T:
+            pred_x_T = pred
+            pred_x_0 = (x_t - B_t * pred_x_T) / A_t
+        elif pred_type == PredictionType.x_0:
+            pred_x_0 = pred
+            pred_x_T = (x_t - A_t * pred_x_0) / B_t
+        elif pred_type == PredictionType.v_cos:
+            pred_x_0 = A_t * x_t - B_t * pred
+            pred_x_T = A_t * pred + B_t * x_t
+        elif pred_type == PredictionType.v_lerp:
+            pred_x_0 = (x_t - B_t * pred) / (A_t + B_t)
+            pred_x_T = (x_t + A_t * pred) / (A_t + B_t)
+        else:
+            raise NotImplementedError
+
+        return pred_x_0, pred_x_T
+
+    def convert_to_pred(
+        self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor, pred_type: PredictionType
+    ) -> torch.FloatTensor:
+        """
+        Convert to prediction target given x_0 and x_T.
+        """
+        if pred_type == PredictionType.x_T:
+            return x_T
+        if pred_type == PredictionType.x_0:
+            return x_0
+        if pred_type == PredictionType.v_cos:
+            t = expand_dims(t, x_0.ndim)
+            return self.A(t) * x_T - self.B(t) * x_0
+        if pred_type == PredictionType.v_lerp:
+            return x_T - x_0
+        raise NotImplementedError

common/diffusion/schedules/lerp.py

@@ -0,0 +1,55 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Linear interpolation schedule (lerp).
+"""
+
+from typing import Union
+import torch
+
+from .base import Schedule
+
+
+class LinearInterpolationSchedule(Schedule):
+    """
+    Linear interpolation schedule (lerp) is proposed by flow matching and rectified flow.
+    It leads to straighter probability flow theoretically. It is also used by Stable Diffusion 3.
+    <https://arxiv.org/abs/2209.03003>
+    <https://arxiv.org/abs/2210.02747>
+
+        x_t = (1 - t) * x_0 + t * x_T
+
+    Can be either continuous or discrete.
+    """
+
+    def __init__(self, T: Union[int, float] = 1.0):
+        self._T = T
+
+    @property
+    def T(self) -> Union[int, float]:
+        return self._T
+
+    def A(self, t: torch.Tensor) -> torch.Tensor:
+        return 1 - (t / self.T)
+
+    def B(self, t: torch.Tensor) -> torch.Tensor:
+        return t / self.T
+
+    # ----------------------------------------------------
+
+    def isnr(self, snr: torch.Tensor) -> torch.Tensor:
+        t = self.T / (1 + snr**0.5)
+        t = t if self.is_continuous() else t.round().int()
+        return t

common/diffusion/timesteps/base.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+from typing import Sequence, Union
+import torch
+
+from ..types import SamplingDirection
+
+
+class Timesteps(ABC):
+    """
+    Timesteps base class.
+    """
+
+    def __init__(self, T: Union[int, float]):
+        assert T > 0
+        self._T = T
+
+    @property
+    def T(self) -> Union[int, float]:
+        """
+        Maximum timestep inclusive.
+        int if discrete, float if continuous.
+        """
+        return self._T
+
+    def is_continuous(self) -> bool:
+        """
+        Whether the schedule is continuous.
+        """
+        return isinstance(self.T, float)
+
+
+class SamplingTimesteps(Timesteps):
+    """
+    Sampling timesteps.
+    It defines the discretization of sampling steps.
+    """
+
+    def __init__(
+        self,
+        T: Union[int, float],
+        timesteps: torch.Tensor,
+        direction: SamplingDirection,
+    ):
+        assert timesteps.ndim == 1
+        super().__init__(T)
+        self.timesteps = timesteps
+        self.direction = direction
+
+    def __len__(self) -> int:
+        """
+        Number of sampling steps.
+        """
+        return len(self.timesteps)
+
+    def __getitem__(self, idx: Union[int, torch.IntTensor]) -> torch.Tensor:
+        """
+        The timestep at the sampling step.
+        Returns a scalar tensor if idx is int,
+        or tensor of the same size if idx is a tensor.
+        """
+        return self.timesteps[idx]
+
+    def index(self, t: torch.Tensor) -> torch.Tensor:
+        """
+        Find index by t.
+        Return index of the same shape as t.
+        Index is -1 if t not found in timesteps.
+        """
+        i, j = t.reshape(-1, 1).eq(self.timesteps).nonzero(as_tuple=True)
+        idx = torch.full_like(t, fill_value=-1, dtype=torch.int)
+        idx.view(-1)[i] = j.int()
+        return idx

common/diffusion/timesteps/sampling/trailing.py

@@ -0,0 +1,49 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+
+from ...types import SamplingDirection
+from ..base import SamplingTimesteps
+
+
+class UniformTrailingSamplingTimesteps(SamplingTimesteps):
+    """
+    Uniform trailing sampling timesteps.
+    Defined in (https://arxiv.org/abs/2305.08891)
+
+    Shift is proposed in SD3 for RF schedule.
+    Defined in (https://arxiv.org/pdf/2403.03206) eq.23
+    """
+
+    def __init__(
+        self,
+        T: int,
+        steps: int,
+        shift: float = 1.0,
+        device: torch.device = "cpu",
+    ):
+        # Create trailing timesteps.
+        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device=device)
+
+        # Shift timesteps.
+        timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)
+
+        # Scale to T range.
+        if isinstance(T, float):
+            timesteps = timesteps * T
+        else:
+            timesteps = timesteps.mul(T + 1).sub(1).round().int()
+
+        super().__init__(T=T, timesteps=timesteps, direction=SamplingDirection.backward)

common/diffusion/types.py

@@ -0,0 +1,59 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Type definitions.
+"""
+
+from enum import Enum
+
+
+class PredictionType(str, Enum):
+    """
+    x_0:
+        Predict data sample.
+    x_T:
+        Predict noise sample.
+        Proposed by DDPM (https://arxiv.org/abs/2006.11239)
+        Proved problematic by zsnr paper (https://arxiv.org/abs/2305.08891)
+    v_cos:
+        Predict velocity dx/dt based on the cosine schedule (A_t * x_T - B_t * x_0).
+        Proposed by progressive distillation (https://arxiv.org/abs/2202.00512)
+    v_lerp:
+        Predict velocity dx/dt based on the lerp schedule (x_T - x_0).
+        Proposed by rectified flow (https://arxiv.org/abs/2209.03003)
+    """
+
+    x_0 = "x_0"
+    x_T = "x_T"
+    v_cos = "v_cos"
+    v_lerp = "v_lerp"
+
+
+class SamplingDirection(str, Enum):
+    """
+    backward: Sample from x_T to x_0 for data generation.
+    forward:  Sample from x_0 to x_T for noise inversion.
+    """
+
+    backward = "backward"
+    forward = "forward"
+
+    @staticmethod
+    def reverse(direction):
+        if direction == SamplingDirection.backward:
+            return SamplingDirection.forward
+        if direction == SamplingDirection.forward:
+            return SamplingDirection.backward
+        raise NotImplementedError

common/diffusion/utils.py

@@ -0,0 +1,84 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Utility functions.
+"""
+
+from typing import Callable
+import torch
+
+
+def expand_dims(tensor: torch.Tensor, ndim: int):
+    """
+    Expand tensor to target ndim. New dims are added to the right.
+    For example, if the tensor shape was (8,), target ndim is 4, return (8, 1, 1, 1).
+    """
+    shape = tensor.shape + (1,) * (ndim - tensor.ndim)
+    return tensor.reshape(shape)
+
+
+def assert_schedule_timesteps_compatible(schedule, timesteps):
+    """
+    Check if schedule and timesteps are compatible.
+    """
+    if schedule.T != timesteps.T:
+        raise ValueError("Schedule and timesteps must have the same T.")
+    if schedule.is_continuous() != timesteps.is_continuous():
+        raise ValueError("Schedule and timesteps must have the same continuity.")
+
+
+def classifier_free_guidance(
+    pos: torch.Tensor,
+    neg: torch.Tensor,
+    scale: float,
+    rescale: float = 0.0,
+):
+    """
+    Apply classifier-free guidance.
+    """
+    # Classifier-free guidance (https://arxiv.org/abs/2207.12598)
+    cfg = neg + scale * (pos - neg)
+
+    # Classifier-free guidance rescale (https://arxiv.org/pdf/2305.08891.pdf)
+    if rescale != 0.0:
+        pos_std = pos.std(dim=list(range(1, pos.ndim)), keepdim=True)
+        cfg_std = cfg.std(dim=list(range(1, cfg.ndim)), keepdim=True)
+        factor = pos_std / cfg_std
+        factor = rescale * factor + (1 - rescale)
+        cfg *= factor
+
+    return cfg
+
+
+def classifier_free_guidance_dispatcher(
+    pos: Callable,
+    neg: Callable,
+    scale: float,
+    rescale: float = 0.0,
+):
+    """
+    Optionally execute models depending on classifer-free guidance scale.
+    """
+    # If scale is 1, no need to execute neg model.
+    if scale == 1.0:
+        return pos()
+
+    # Otherwise, execute both pos nad neg models and apply cfg.
+    return classifier_free_guidance(
+        pos=pos(),
+        neg=neg(),
+        scale=scale,
+        rescale=rescale,
+    )

common/distributed/__init__.py

@@ -0,0 +1,37 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed package.
+"""
+
+from .basic import (
+    barrier_if_distributed,
+    convert_to_ddp,
+    get_device,
+    get_global_rank,
+    get_local_rank,
+    get_world_size,
+    init_torch,
+)
+
+__all__ = [
+    "barrier_if_distributed",
+    "convert_to_ddp",
+    "get_device",
+    "get_global_rank",
+    "get_local_rank",
+    "get_world_size",
+    "init_torch",
+]

common/distributed/advanced.py

@@ -0,0 +1,208 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Advanced distributed functions for sequence parallel.
+"""
+
+from typing import Optional, List
+import torch
+import torch.distributed as dist
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.distributed.fsdp import ShardingStrategy
+
+from .basic import get_global_rank, get_world_size
+
+
+_DATA_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_GROUP = None
+_SEQUENCE_PARALLEL_CPU_GROUP = None
+_MODEL_SHARD_CPU_INTER_GROUP = None
+_MODEL_SHARD_CPU_INTRA_GROUP = None
+_MODEL_SHARD_INTER_GROUP = None
+_MODEL_SHARD_INTRA_GROUP = None
+_SEQUENCE_PARALLEL_GLOBAL_RANKS = None
+
+
+def get_data_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get data parallel process group.
+    """
+    return _DATA_PARALLEL_GROUP
+
+
+def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel process group.
+    """
+    return _SEQUENCE_PARALLEL_GROUP
+
+
+def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get sequence parallel CPU process group.
+    """
+    return _SEQUENCE_PARALLEL_CPU_GROUP
+
+
+def get_data_parallel_rank() -> int:
+    """
+    Get data parallel rank.
+    """
+    group = get_data_parallel_group()
+    return dist.get_rank(group) if group else get_global_rank()
+
+
+def get_data_parallel_world_size() -> int:
+    """
+    Get data parallel world size.
+    """
+    group = get_data_parallel_group()
+    return dist.get_world_size(group) if group else get_world_size()
+
+
+def get_sequence_parallel_rank() -> int:
+    """
+    Get sequence parallel rank.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_rank(group) if group else 0
+
+
+def get_sequence_parallel_world_size() -> int:
+    """
+    Get sequence parallel world size.
+    """
+    group = get_sequence_parallel_group()
+    return dist.get_world_size(group) if group else 1
+
+
+def get_model_shard_cpu_intra_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the CPU intra process group of model sharding.
+    """
+    return _MODEL_SHARD_CPU_INTRA_GROUP
+
+
+def get_model_shard_cpu_inter_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the CPU inter process group of model sharding.
+    """
+    return _MODEL_SHARD_CPU_INTER_GROUP
+
+
+def get_model_shard_intra_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the GPU intra process group of model sharding.
+    """
+    return _MODEL_SHARD_INTRA_GROUP
+
+
+def get_model_shard_inter_group() -> Optional[dist.ProcessGroup]:
+    """
+    Get the GPU inter process group of model sharding.
+    """
+    return _MODEL_SHARD_INTER_GROUP
+
+
+def init_sequence_parallel(sequence_parallel_size: int):
+    """
+    Initialize sequence parallel.
+    """
+    global _DATA_PARALLEL_GROUP
+    global _SEQUENCE_PARALLEL_GROUP
+    global _SEQUENCE_PARALLEL_CPU_GROUP
+    global _SEQUENCE_PARALLEL_GLOBAL_RANKS
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    data_parallel_size = world_size // sequence_parallel_size
+    for i in range(data_parallel_size):
+        start_rank = i * sequence_parallel_size
+        end_rank = (i + 1) * sequence_parallel_size
+        ranks = range(start_rank, end_rank)
+        group = dist.new_group(ranks)
+        cpu_group = dist.new_group(ranks, backend="gloo")
+        if rank in ranks:
+            _SEQUENCE_PARALLEL_GROUP = group
+            _SEQUENCE_PARALLEL_CPU_GROUP = cpu_group
+            _SEQUENCE_PARALLEL_GLOBAL_RANKS = list(ranks)
+
+
+def init_model_shard_group(
+    *,
+    sharding_strategy: ShardingStrategy,
+    device_mesh: Optional[DeviceMesh] = None,
+):
+    """
+    Initialize process group of model sharding.
+    """
+    global _MODEL_SHARD_INTER_GROUP
+    global _MODEL_SHARD_INTRA_GROUP
+    global _MODEL_SHARD_CPU_INTER_GROUP
+    global _MODEL_SHARD_CPU_INTRA_GROUP
+    assert dist.is_initialized()
+    world_size = dist.get_world_size()
+    if device_mesh is not None:
+        num_shards_per_group = device_mesh.shape[1]
+    elif sharding_strategy == ShardingStrategy.NO_SHARD:
+        num_shards_per_group = 1
+    elif sharding_strategy in [
+        ShardingStrategy.HYBRID_SHARD,
+        ShardingStrategy._HYBRID_SHARD_ZERO2,
+    ]:
+        num_shards_per_group = torch.cuda.device_count()
+    else:
+        num_shards_per_group = world_size
+    num_groups = world_size // num_shards_per_group
+    device_mesh = (num_groups, num_shards_per_group)
+
+    gpu_mesh_2d = init_device_mesh("cuda", device_mesh, mesh_dim_names=("inter", "intra"))
+    cpu_mesh_2d = init_device_mesh("cpu", device_mesh, mesh_dim_names=("inter", "intra"))
+
+    _MODEL_SHARD_INTER_GROUP = gpu_mesh_2d.get_group("inter")
+    _MODEL_SHARD_INTRA_GROUP = gpu_mesh_2d.get_group("intra")
+    _MODEL_SHARD_CPU_INTER_GROUP = cpu_mesh_2d.get_group("inter")
+    _MODEL_SHARD_CPU_INTRA_GROUP = cpu_mesh_2d.get_group("intra")
+
+def get_sequence_parallel_global_ranks() -> List[int]:
+    """
+    Get all global ranks of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    if _SEQUENCE_PARALLEL_GLOBAL_RANKS is None:
+        return [dist.get_rank()]
+    return _SEQUENCE_PARALLEL_GLOBAL_RANKS
+
+
+def get_next_sequence_parallel_rank() -> int:
+    """
+    Get the next global rank of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    sp_global_ranks = get_sequence_parallel_global_ranks()
+    sp_rank = get_sequence_parallel_rank()
+    sp_size = get_sequence_parallel_world_size()
+    return sp_global_ranks[(sp_rank + 1) % sp_size]
+
+
+def get_prev_sequence_parallel_rank() -> int:
+    """
+    Get the previous global rank of the sequence parallel process group
+    that the caller rank belongs to.
+    """
+    sp_global_ranks = get_sequence_parallel_global_ranks()
+    sp_rank = get_sequence_parallel_rank()
+    sp_size = get_sequence_parallel_world_size()
+    return sp_global_ranks[(sp_rank + sp_size - 1) % sp_size]

common/distributed/basic.py

@@ -0,0 +1,84 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed basic functions.
+"""
+
+import os
+from datetime import timedelta
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+
+def get_global_rank() -> int:
+    """
+    Get the global rank, the global index of the GPU.
+    """
+    return int(os.environ.get("RANK", "0"))
+
+
+def get_local_rank() -> int:
+    """
+    Get the local rank, the local index of the GPU.
+    """
+    return int(os.environ.get("LOCAL_RANK", "0"))
+
+
+def get_world_size() -> int:
+    """
+    Get the world size, the total amount of GPUs.
+    """
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def get_device() -> torch.device:
+    """
+    Get current rank device.
+    """
+    return torch.device("cuda", get_local_rank())
+
+
+def barrier_if_distributed(*args, **kwargs):
+    """
+    Synchronizes all processes if under distributed context.
+    """
+    if dist.is_initialized():
+        return dist.barrier(*args, **kwargs)
+
+
+def init_torch(cudnn_benchmark=True, timeout=timedelta(seconds=600)):
+    """
+    Common PyTorch initialization configuration.
+    """
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = cudnn_benchmark
+    torch.cuda.set_device(get_local_rank())
+    dist.init_process_group(
+        backend="nccl",
+        rank=get_global_rank(),
+        world_size=get_world_size(),
+        timeout=timeout,
+    )
+
+
+def convert_to_ddp(module: torch.nn.Module, **kwargs) -> DistributedDataParallel:
+    return DistributedDataParallel(
+        module=module,
+        device_ids=[get_local_rank()],
+        output_device=get_local_rank(),
+        **kwargs,
+    )

common/distributed/meta_init_utils.py

@@ -0,0 +1,41 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import torch
+from rotary_embedding_torch import RotaryEmbedding
+from torch import nn
+from torch.distributed.fsdp._common_utils import _is_fsdp_flattened
+
+__all__ = ["meta_non_persistent_buffer_init_fn"]
+
+
+def meta_non_persistent_buffer_init_fn(module: nn.Module) -> nn.Module:
+    """
+    Used for materializing `non-persistent tensor buffers` while model resuming.
+
+    Since non-persistent tensor buffers are not saved in state_dict,
+    when initializing model with meta device, user should materialize those buffers manually.
+
+    Currently, only `rope.dummy` is this special case.
+    """
+    with torch.no_grad():
+        for submodule in module.modules():
+            if not isinstance(submodule, RotaryEmbedding):
+                continue
+            for buffer_name, buffer in submodule.named_buffers(recurse=False):
+                if buffer.is_meta and "dummy" in buffer_name:
+                    materialized_buffer = torch.zeros_like(buffer, device="cpu")
+                    setattr(submodule, buffer_name, materialized_buffer)
+    assert not any(b.is_meta for n, b in module.named_buffers())
+    return module

common/distributed/ops.py

@@ -0,0 +1,494 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Distributed ops for supporting sequence parallel.
+"""
+
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+from common.cache import Cache
+from common.distributed.advanced import (
+    get_sequence_parallel_group,
+    get_sequence_parallel_rank,
+    get_sequence_parallel_world_size,
+)
+
+from .basic import get_device
+
+_SEQ_DATA_BUF = defaultdict(lambda: [None, None, None])
+_SEQ_DATA_META_SHAPES = defaultdict()
+_SEQ_DATA_META_DTYPES = defaultdict()
+_SEQ_DATA_ASYNC_COMMS = defaultdict(list)
+_SYNC_BUFFER = defaultdict(dict)
+
+
+def single_all_to_all(
+    local_input: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+    async_op: bool = False,
+):
+    """
+    A function to do all-to-all on a tensor
+    """
+    seq_world_size = dist.get_world_size(group)
+    prev_scatter_dim = scatter_dim
+    if scatter_dim != 0:
+        local_input = local_input.transpose(0, scatter_dim)
+        if gather_dim == 0:
+            gather_dim = scatter_dim
+        scatter_dim = 0
+
+    inp_shape = list(local_input.shape)
+    inp_shape[scatter_dim] = inp_shape[scatter_dim] // seq_world_size
+    input_t = local_input.reshape(
+        [seq_world_size, inp_shape[scatter_dim]] + inp_shape[scatter_dim + 1 :]
+    ).contiguous()
+    output = torch.empty_like(input_t)
+    comm = dist.all_to_all_single(output, input_t, group=group, async_op=async_op)
+    if async_op:
+        # let user's code transpose & reshape
+        return output, comm, prev_scatter_dim
+
+    # first dim is seq_world_size, so we can split it directly
+    output = torch.cat(output.split(1), dim=gather_dim + 1).squeeze(0)
+    if prev_scatter_dim:
+        output = output.transpose(0, prev_scatter_dim).contiguous()
+    return output
+
+
+def _all_to_all(
+    local_input: Tensor,
+    scatter_dim: int,
+    gather_dim: int,
+    group: dist.ProcessGroup,
+):
+    seq_world_size = dist.get_world_size(group)
+    input_list = [
+        t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)
+    ]
+    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]
+    dist.all_to_all(output_list, input_list, group=group)
+    return torch.cat(output_list, dim=gather_dim).contiguous()
+
+
+class SeqAllToAll(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        scatter_dim: int,
+        gather_dim: int,
+        async_op: bool,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.scatter_dim = scatter_dim
+        ctx.gather_dim = gather_dim
+        ctx.async_op = async_op
+        if async_op:
+            output, comm, prev_scatter_dim = single_all_to_all(
+                local_input, scatter_dim, gather_dim, group, async_op=async_op
+            )
+            ctx.prev_scatter_dim = prev_scatter_dim
+            return output, comm
+
+        return _all_to_all(local_input, scatter_dim, gather_dim, group)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]:
+        if ctx.async_op:
+            input_t = torch.cat(grad_output[0].split(1), dim=ctx.gather_dim + 1).squeeze(0)
+            if ctx.prev_scatter_dim:
+                input_t = input_t.transpose(0, ctx.prev_scatter_dim)
+        else:
+            input_t = grad_output[0]
+        return (
+            None,
+            _all_to_all(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group),
+            None,
+            None,
+            None,
+        )
+
+
+class Slice(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        ctx.dim = dim
+        dim_size = local_input.shape[dim]
+        return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous()
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]:
+        dim_size = list(grad_output.size())
+        split_size = dim_size[0]
+        dim_size[0] = dim_size[0] * ctx.seq_world_size
+        output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device())
+        dist._all_gather_base(output, grad_output, group=ctx.group)
+        return (None, torch.cat(output.split(split_size), dim=ctx.dim), None)
+
+
+class Gather(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        group: dist.ProcessGroup,
+        local_input: Tensor,
+        dim: int,
+        grad_scale: Optional[bool] = False,
+    ) -> Tensor:
+        ctx.group = group
+        ctx.rank = dist.get_rank(group)
+        ctx.dim = dim
+        ctx.grad_scale = grad_scale
+        seq_world_size = dist.get_world_size(group)
+        ctx.seq_world_size = seq_world_size
+        dim_size = list(local_input.size())
+        split_size = dim_size[0]
+        ctx.part_size = dim_size[dim]
+        dim_size[0] = dim_size[0] * seq_world_size
+        output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device())
+        dist._all_gather_base(output, local_input.contiguous(), group=ctx.group)
+        return torch.cat(output.split(split_size), dim=dim)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]:
+        if ctx.grad_scale:
+            grad_output = grad_output * ctx.seq_world_size
+        return (
+            None,
+            grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(),
+            None,
+            None,
+        )
+
+
+def gather_seq_scatter_heads_qkv(
+    qkv_tensor: Tensor,
+    *,
+    seq_dim: int,
+    qkv_shape: Optional[Tensor] = None,
+    cache: Cache = Cache(disable=True),
+    restore_shape: bool = True,
+):
+    """
+    A func to sync splited qkv tensor
+    qkv_tensor: the tensor we want to do alltoall with. The last dim must
+        be the projection_idx, which we will split into 3 part. After
+        spliting, the gather idx will be projecttion_idx + 1
+    seq_dim: gather_dim for all2all comm
+    restore_shape: if True, output will has the same shape length as input
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return qkv_tensor
+    world = get_sequence_parallel_world_size()
+    orig_shape = qkv_tensor.shape
+    scatter_dim = qkv_tensor.dim()
+    bef_all2all_shape = list(orig_shape)
+    qkv_proj_dim = bef_all2all_shape[-1]
+    bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3]
+    qkv_tensor = qkv_tensor.view(bef_all2all_shape)
+    qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, False)
+    if restore_shape:
+        out_shape = list(orig_shape)
+        out_shape[seq_dim] *= world
+        out_shape[-1] = qkv_proj_dim // world
+        qkv_tensor = qkv_tensor.view(out_shape)
+
+    # remove padding
+    if qkv_shape is not None:
+        unpad_dim_size = cache(
+            "unpad_dim_size", lambda: torch.sum(torch.prod(qkv_shape, dim=-1)).item()
+        )
+        if unpad_dim_size % world != 0:
+            padding_size = qkv_tensor.size(seq_dim) - unpad_dim_size
+            qkv_tensor = _unpad_tensor(qkv_tensor, seq_dim, padding_size)
+    return qkv_tensor
+
+
+def slice_inputs(x: Tensor, dim: int, padding: bool = True):
+    """
+    A func to slice the input sequence in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if group is None:
+        return x
+    sp_rank = get_sequence_parallel_rank()
+    sp_world = get_sequence_parallel_world_size()
+    dim_size = x.shape[dim]
+    unit = (dim_size + sp_world - 1) // sp_world
+    if padding and dim_size % sp_world:
+        padding_size = sp_world - (dim_size % sp_world)
+        x = _pad_tensor(x, dim, padding_size)
+    slc = [slice(None)] * len(x.shape)
+    slc[dim] = slice(unit * sp_rank, unit * (sp_rank + 1))
+    return x[slc]
+
+
+def remove_seqeunce_parallel_padding(x: Tensor, dim: int, unpad_dim_size: int):
+    """
+    A func to remove the padding part of the tensor based on its original shape
+    """
+    group = get_sequence_parallel_group()
+    if group is None:
+        return x
+    sp_world = get_sequence_parallel_world_size()
+    if unpad_dim_size % sp_world == 0:
+        return x
+    padding_size = sp_world - (unpad_dim_size % sp_world)
+    assert (padding_size + unpad_dim_size) % sp_world == 0
+    return _unpad_tensor(x, dim=dim, padding_size=padding_size)
+
+
+def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor:
+    """
+    A func to sync attention result with alltoall in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    dim_size = x.size(seq_dim)
+    sp_world = get_sequence_parallel_world_size()
+    if dim_size % sp_world != 0:
+        padding_size = sp_world - (dim_size % sp_world)
+        x = _pad_tensor(x, seq_dim, padding_size)
+    return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)
+
+
+def gather_seq_scatter_heads(x: Tensor, seq_dim: int, head_dim: int) -> Tensor:
+    """
+    A func to sync embedding input with alltoall in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return SeqAllToAll.apply(group, x, head_dim, seq_dim, False)
+
+
+def scatter_heads(x: Tensor, dim: int) -> Tensor:
+    """
+    A func to split heads before attention in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return Slice.apply(group, x, dim)
+
+
+def gather_heads(x: Tensor, dim: int, grad_scale: Optional[bool] = False) -> Tensor:
+    """
+    A func to gather heads for the attention result in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    return Gather.apply(group, x, dim, grad_scale)
+
+
+def gather_outputs(
+    x: Tensor,
+    *,
+    gather_dim: int,
+    padding_dim: Optional[int] = None,
+    unpad_shape: Optional[Tensor] = None,
+    cache: Cache = Cache(disable=True),
+    scale_grad=True,
+):
+    """
+    A func to gather the outputs for the model result in sequence parallel
+    """
+    group = get_sequence_parallel_group()
+    if not group:
+        return x
+    x = Gather.apply(group, x, gather_dim, scale_grad)
+    if padding_dim is not None:
+        unpad_dim_size = cache(
+            "unpad_dim_size", lambda: torch.sum(torch.prod(unpad_shape, dim=1)).item()
+        )
+        x = remove_seqeunce_parallel_padding(x, padding_dim, unpad_dim_size)
+    return x
+
+
+def _pad_tensor(x: Tensor, dim: int, padding_size: int):
+    shape = list(x.shape)
+    shape[dim] = padding_size
+    pad = torch.zeros(shape, dtype=x.dtype, device=x.device)
+    return torch.cat([x, pad], dim=dim)
+
+
+def _unpad_tensor(x: Tensor, dim: int, padding_size):
+    slc = [slice(None)] * len(x.shape)
+    slc[dim] = slice(0, -padding_size)
+    return x[slc]
+
+
+def _broadcast_data(data, shape, dtype, src, group, async_op):
+    comms = []
+    if isinstance(data, (list, tuple)):
+        for i, sub_shape in enumerate(shape):
+            comms += _broadcast_data(data[i], sub_shape, dtype[i], src, group, async_op)
+    elif isinstance(data, dict):
+        for key, sub_data in data.items():
+            comms += _broadcast_data(sub_data, shape[key], dtype[key], src, group, async_op)
+    elif isinstance(data, Tensor):
+        comms.append(dist.broadcast(data, src=src, group=group, async_op=async_op))
+    return comms
+
+
+def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]:
+    if isinstance(data, (list, tuple)):
+        return [_traverse(sub_data, op) for sub_data in data]
+    elif isinstance(data, dict):
+        return {key: _traverse(sub_data, op) for key, sub_data in data.items()}
+    elif isinstance(data, Tensor):
+        return op(data)
+    else:
+        return None
+
+
+def _get_shapes(data):
+    return _traverse(data, op=lambda x: x.shape)
+
+
+def _get_dtypes(data):
+    return _traverse(data, op=lambda x: x.dtype)
+
+
+def _construct_broadcast_buffer(shapes, dtypes, device):
+    if isinstance(shapes, torch.Size):
+        return torch.empty(shapes, dtype=dtypes, device=device)
+
+    if isinstance(shapes, (list, tuple)):
+        buffer = []
+        for i, sub_shape in enumerate(shapes):
+            buffer.append(_construct_broadcast_buffer(sub_shape, dtypes[i], device))
+    elif isinstance(shapes, dict):
+        buffer = {}
+        for key, sub_shape in shapes.items():
+            buffer[key] = _construct_broadcast_buffer(sub_shape, dtypes[key], device)
+    else:
+        return None
+    return buffer
+
+
+class SPDistForward:
+    """A forward tool to sync different result across sp group
+
+    Args:
+        module: a function or module to process users input
+        sp_step: current training step to judge which rank to broadcast its result to all
+        name: a distinct str to save meta and async comm
+        comm_shape: if different ranks have different shape, mark this arg to True
+        device: the device for current rank, can be empty
+    """
+
+    def __init__(
+        self,
+        name: str,
+        comm_shape: bool,
+        device: torch.device = None,
+    ):
+        self.name = name
+        self.comm_shape = comm_shape
+        if device:
+            self.device = device
+        else:
+            self.device = get_device()
+
+    def __call__(self, inputs) -> Any:
+        group = get_sequence_parallel_group()
+        if not group:
+            yield inputs
+        else:
+            device = self.device
+            sp_world = get_sequence_parallel_world_size()
+            sp_rank = get_sequence_parallel_rank()
+            for local_step in range(sp_world):
+                src_rank = dist.get_global_rank(group, local_step)
+                is_src = sp_rank == local_step
+                local_shapes = []
+                local_dtypes = []
+                if local_step == 0:
+                    local_result = inputs
+                    _SEQ_DATA_BUF[self.name][-1] = local_result
+                    local_shapes = _get_shapes(local_result)
+                    local_dtypes = _get_dtypes(local_result)
+                    if self.comm_shape:
+                        group_shapes_lists = [None] * sp_world
+                        dist.all_gather_object(group_shapes_lists, local_shapes, group=group)
+                        _SEQ_DATA_META_SHAPES[self.name] = group_shapes_lists
+                    else:
+                        _SEQ_DATA_META_SHAPES[self.name] = [local_shapes] * sp_world
+                    _SEQ_DATA_META_DTYPES[self.name] = local_dtypes
+                shapes = _SEQ_DATA_META_SHAPES[self.name][local_step]
+                dtypes = _SEQ_DATA_META_DTYPES[self.name]
+                buf_id = local_step % 2
+                if local_step == 0:
+                    sync_data = (
+                        local_result
+                        if is_src
+                        else _construct_broadcast_buffer(shapes, dtypes, device)
+                    )
+                    _broadcast_data(sync_data, shapes, dtypes, src_rank, group, False)
+                    _SEQ_DATA_BUF[self.name][buf_id] = sync_data
+
+                # wait for async comm ops
+                if _SEQ_DATA_ASYNC_COMMS[self.name]:
+                    for comm in _SEQ_DATA_ASYNC_COMMS[self.name]:
+                        comm.wait()
+                # before return the sync result, do async broadcast for next batch
+                if local_step < sp_world - 1:
+                    next_buf_id = 1 - buf_id
+                    shapes = _SEQ_DATA_META_SHAPES[self.name][local_step + 1]
+                    src_rank = dist.get_global_rank(group, local_step + 1)
+                    is_src = sp_rank == local_step + 1
+                    next_sync_data = (
+                        _SEQ_DATA_BUF[self.name][-1]
+                        if is_src
+                        else _construct_broadcast_buffer(shapes, dtypes, device)
+                    )
+                    _SEQ_DATA_ASYNC_COMMS[self.name] = _broadcast_data(
+                        next_sync_data, shapes, dtypes, src_rank, group, True
+                    )
+                    _SEQ_DATA_BUF[self.name][next_buf_id] = next_sync_data
+                yield _SEQ_DATA_BUF[self.name][buf_id]
+
+
+sync_inputs = SPDistForward(name="bef_fwd", comm_shape=True)
+
+
+def sync_data(data, sp_idx, name="tmp"):
+    group = get_sequence_parallel_group()
+    if group is None:
+        return data
+    # if sp_idx in _SYNC_BUFFER[name]:
+    #     return _SYNC_BUFFER[name][sp_idx]
+    sp_rank = get_sequence_parallel_rank()
+    src_rank = dist.get_global_rank(group, sp_idx)
+    objects = [data] if sp_rank == sp_idx else [None]
+    dist.broadcast_object_list(objects, src=src_rank, group=group)
+    # _SYNC_BUFFER[name] = {sp_idx: objects[0]}
+    return objects[0]

common/logger.py

@@ -0,0 +1,44 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Logging utility functions.
+"""
+
+import logging
+import sys
+from typing import Optional
+
+from common.distributed import get_global_rank, get_local_rank, get_world_size
+
+_default_handler = logging.StreamHandler(sys.stdout)
+_default_handler.setFormatter(
+    logging.Formatter(
+        "%(asctime)s "
+        + (f"[Rank:{get_global_rank()}]" if get_world_size() > 1 else "")
+        + (f"[LocalRank:{get_local_rank()}]" if get_world_size() > 1 else "")
+        + "[%(threadName).12s][%(name)s][%(levelname).5s] "
+        + "%(message)s"
+    )
+)
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    """
+    Get a logger.
+    """
+    logger = logging.getLogger(name)
+    logger.addHandler(_default_handler)
+    logger.setLevel(logging.INFO)
+    return logger

common/partition.py

@@ -0,0 +1,59 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+"""
+Partition utility functions.
+"""
+
+from typing import Any, List
+
+
+def partition_by_size(data: List[Any], size: int) -> List[List[Any]]:
+    """
+    Partition a list by size.
+    When indivisible, the last group contains fewer items than the target size.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - size: 2
+        - return: [[1,2], [3,4], [5]]
+    """
+    assert size > 0
+    return [data[i : (i + size)] for i in range(0, len(data), size)]
+
+
+def partition_by_groups(data: List[Any], groups: int) -> List[List[Any]]:
+    """
+    Partition a list by groups.
+    When indivisible, some groups may have more items than others.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - groups: 2
+        - return: [[1,3,5], [2,4]]
+    """
+    assert groups > 0
+    return [data[i::groups] for i in range(groups)]
+
+
+def shift_list(data: List[Any], n: int) -> List[Any]:
+    """
+    Rotate a list by n elements.
+
+    Examples:
+        - data: [1,2,3,4,5]
+        - n: 3
+        - return: [4,5,1,2,3]
+    """
+    return data[(n % len(data)) :] + data[: (n % len(data))]

common/seed.py

@@ -0,0 +1,30 @@
+# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# //
+# // Licensed under the Apache License, Version 2.0 (the "License");
+# // you may not use this file except in compliance with the License.
+# // You may obtain a copy of the License at
+# //
+# //     http://www.apache.org/licenses/LICENSE-2.0
+# //
+# // Unless required by applicable law or agreed to in writing, software
+# // distributed under the License is distributed on an "AS IS" BASIS,
+# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# // See the License for the specific language governing permissions and
+# // limitations under the License.
+
+import random
+from typing import Optional
+import numpy as np
+import torch
+
+from common.distributed import get_global_rank
+
+
+def set_seed(seed: Optional[int], same_across_ranks: bool = False):
+    """Function that sets the seed for pseudo-random number generators."""
+    if seed is not None:
+        seed += get_global_rank() if not same_across_ranks else 0
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+

common/utils.py

@@ -0,0 +1,232 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import re
+
+import cv2
+import numpy as np
+import torch
+from torchvision.utils import make_grid
+
+
+# from basicsr
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    return _totensor(imgs, bgr2rgb, float32)
+
+
+def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
+    """Convert torch Tensors into image numpy arrays.
+
+    After clamping to [min, max], values will be normalized to [0, 1].
+
+    Args:
+        tensor (Tensor or list[Tensor]): Accept shapes:
+            1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
+            2) 3D Tensor of shape (3/1 x H x W);
+            3) 2D Tensor of shape (H x W).
+            Tensor channel should be in RGB order.
+        rgb2bgr (bool): Whether to change rgb to bgr.
+        out_type (numpy type): output types. If ``np.uint8``, transform outputs
+            to uint8 type with range [0, 255]; otherwise, float type with
+            range [0, 1]. Default: ``np.uint8``.
+        min_max (tuple[int]): min and max values for clamp.
+
+    Returns:
+        (Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
+        shape (H x W). The channel order is BGR.
+    """
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
+
+    if torch.is_tensor(tensor):
+        tensor = [tensor]
+    result = []
+    for _tensor in tensor:
+        _tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
+        _tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
+
+        n_dim = _tensor.dim()
+        if n_dim == 4:
+            img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if rgb2bgr:
+                img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 3:
+            img_np = _tensor.numpy()
+            img_np = img_np.transpose(1, 2, 0)
+            if img_np.shape[2] == 1:  # gray image
+                img_np = np.squeeze(img_np, axis=2)
+            else:
+                if rgb2bgr:
+                    img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+        elif n_dim == 2:
+            img_np = _tensor.numpy()
+        else:
+            raise TypeError(f'Only support 4D, 3D or 2D tensor. But received with dimension: {n_dim}')
+        if out_type == np.uint8:
+            # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
+            img_np = (img_np * 255.0).round()
+        img_np = img_np.astype(out_type)
+        result.append(img_np)
+    if len(result) == 1:
+        result = result[0]
+    return result
+
+
+def resize_numpy_image_area(image, area=512 * 512):
+    h, w = image.shape[:2]
+    k = math.sqrt(area / (h * w))
+    h = int(h * k) - (int(h * k) % 16)
+    w = int(w * k) - (int(w * k) % 16)
+    image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
+    return image
+
+def resize_numpy_image_long(image, long_edge=768):
+    h, w = image.shape[:2]
+    if max(h, w) <= long_edge:
+        return image
+    k = long_edge / max(h, w)
+    h = int(h * k)
+    w = int(w * k)
+    image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
+    return image
+
+
+# reference: https://github.com/huggingface/diffusers/pull/9295/files
+def convert_flux_lora_to_diffusers(old_state_dict):
+    new_state_dict = {}
+    orig_keys = list(old_state_dict.keys())
+
+    def handle_qkv(sds_sd, ait_sd, sds_key, ait_keys, dims=None):
+        down_weight = sds_sd.pop(sds_key)
+        up_weight = sds_sd.pop(sds_key.replace(".down.weight", ".up.weight"))
+
+        # calculate dims if not provided
+        num_splits = len(ait_keys)
+        if dims is None:
+            dims = [up_weight.shape[0] // num_splits] * num_splits
+        else:
+            assert sum(dims) == up_weight.shape[0]
+
+        # make ai-toolkit weight
+        ait_down_keys = [k + ".lora_A.weight" for k in ait_keys]
+        ait_up_keys = [k + ".lora_B.weight" for k in ait_keys]
+
+        # down_weight is copied to each split
+        ait_sd.update({k: down_weight for k in ait_down_keys})
+
+        # up_weight is split to each split
+        ait_sd.update({k: v for k, v in zip(ait_up_keys, torch.split(up_weight, dims, dim=0))})  # noqa: C416
+
+    for old_key in orig_keys:
+        # Handle double_blocks
+        if 'double_blocks' in old_key:
+            block_num = re.search(r"double_blocks_(\d+)", old_key).group(1)
+            new_key = f"transformer.transformer_blocks.{block_num}"
+
+            if "proj_lora1" in old_key:
+                new_key += ".attn.to_out.0"
+            elif "proj_lora2" in old_key:
+                new_key += ".attn.to_add_out"
+            elif "qkv_lora2" in old_key and "up" not in old_key:
+                handle_qkv(
+                    old_state_dict,
+                    new_state_dict,
+                    old_key,
+                    [
+                        f"transformer.transformer_blocks.{block_num}.attn.add_q_proj",
+                        f"transformer.transformer_blocks.{block_num}.attn.add_k_proj",
+                        f"transformer.transformer_blocks.{block_num}.attn.add_v_proj",
+                    ],
+                )
+                # continue
+            elif "qkv_lora1" in old_key and "up" not in old_key:
+                handle_qkv(
+                    old_state_dict,
+                    new_state_dict,
+                    old_key,
+                    [
+                        f"transformer.transformer_blocks.{block_num}.attn.to_q",
+                        f"transformer.transformer_blocks.{block_num}.attn.to_k",
+                        f"transformer.transformer_blocks.{block_num}.attn.to_v",
+                    ],
+                )
+                # continue
+
+            if "down" in old_key:
+                new_key += ".lora_A.weight"
+            elif "up" in old_key:
+                new_key += ".lora_B.weight"
+
+        # Handle single_blocks
+        elif 'single_blocks' in old_key:
+            block_num = re.search(r"single_blocks_(\d+)", old_key).group(1)
+            new_key = f"transformer.single_transformer_blocks.{block_num}"
+
+            if "proj_lora" in old_key:
+                new_key += ".proj_out"
+            elif "qkv_lora" in old_key and "up" not in old_key:
+                handle_qkv(
+                    old_state_dict,
+                    new_state_dict,
+                    old_key,
+                    [
+                        f"transformer.single_transformer_blocks.{block_num}.attn.to_q",
+                        f"transformer.single_transformer_blocks.{block_num}.attn.to_k",
+                        f"transformer.single_transformer_blocks.{block_num}.attn.to_v",
+                    ],
+                )
+
+            if "down" in old_key:
+                new_key += ".lora_A.weight"
+            elif "up" in old_key:
+                new_key += ".lora_B.weight"
+
+        else:
+            # Handle other potential key patterns here
+            new_key = old_key
+
+        # Since we already handle qkv above.
+        if "qkv" not in old_key and 'embedding' not in old_key:
+            new_state_dict[new_key] = old_state_dict.pop(old_key)
+
+    # if len(old_state_dict) > 0:
+    #     raise ValueError(f"`old_state_dict` should be at this point but has: {list(old_state_dict.keys())}.")
+
+    return new_state_dict