models/dreamgen/examples/text2image.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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 argparse
import json
import os

from imaginaire.auxiliary.text_encoder import CosmosTextEncoder
from imaginaire.lazy_config.lazy import LazyConfig

# Set TOKENIZERS_PARALLELISM environment variable to avoid deadlocks with multiprocessing
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import time

import torch
from megatron.core import parallel_state

from cosmos_predict2.configs.base.config_text2image import (
    get_cosmos_predict2_text2image_pipeline,
)
from cosmos_predict2.pipelines.text2image import Text2ImagePipeline
from imaginaire.constants import (
    CosmosPredict2Text2ImageModelSize,
    CosmosPredict2Video2WorldAspectRatio,
    get_cosmos_predict2_text2image_checkpoint,
    print_environment_info,
)
from imaginaire.utils import distributed, log, misc
from imaginaire.utils.io import save_image_or_video, save_text_prompts

# ---- Sparse Backend Integration ----
def _import_sparse_backends():
    import sys
    # Search for project root (where 'methods' directory resides)
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # The script is in models/dreamgen/examples/
    # Project root should be 3 levels up
    project_root = os.path.normpath(os.path.join(current_dir, "../../.."))
    if project_root not in sys.path:
        sys.path.insert(0, project_root)

    from methods.cache_strategy.FasterCache.config import add_fastercache_args
    from methods.cache_strategy.FasterCache.runtime import apply_fastercache
    from methods.prunning.SiTo.config import add_dreamgen_sito_args
    from methods.prunning.importance_token_merge.config import add_dreamgen_itm_args
    from methods.sparse_attention.pisa.config import add_pisa_args
    from methods.sparse_attention.svg.config import add_dreamgen_svg_args
    from methods.sparse_attention.radial.config import add_dreamgen_radial_args
    return (
        add_pisa_args,
        add_dreamgen_svg_args,
        add_dreamgen_radial_args,
        add_dreamgen_sito_args,
        add_dreamgen_itm_args,
        add_fastercache_args,
        apply_fastercache,
    )
# -----------------------------------

_DEFAULT_POSITIVE_PROMPT = "A well-worn broom sweeps across a dusty wooden floor, its bristles gathering crumbs and flecks of debris in swift, rhythmic strokes. Dust motes dance in the sunbeams filtering through the window, glowing momentarily before settling. The quiet swish of straw brushing wood is interrupted only by the occasional creak of old floorboards. With each pass, the floor grows cleaner, restoring a sense of quiet order to the humble room."


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Text to Image Generation with Cosmos Predict2")
    parser.add_argument(
        "--model_size",
        choices=CosmosPredict2Text2ImageModelSize.__args__,
        default="2B",
        help="Size of the model to use for text-to-image generation",
    )
    parser.add_argument(
        "--distill_steps",
        type=int,
        choices=[0, 1, 2, 3, 4],
        default=0,
        help="1~4 for timestep-distilled inference; 0 for the original non-distilled model",
    )
    parser.add_argument(
        "--dit_path",
        type=str,
        default="",
        help="Custom path to the DiT model checkpoint for post-trained models.",
    )
    parser.add_argument(
        "--load_ema",
        action="store_true",
        help="Use EMA weights for generation.",
    )
    parser.add_argument("--prompt", type=str, default=_DEFAULT_POSITIVE_PROMPT, help="Text prompt for image generation")
    parser.add_argument(
        "--batch_input_json",
        type=str,
        default=None,
        help="Path to JSON file containing batch inputs. Each entry should have 'prompt' and 'output_image' fields.",
    )
    parser.add_argument("--negative_prompt", type=str, default="", help="Negative text prompt for image generation")
    parser.add_argument(
        "--aspect_ratio",
        choices=CosmosPredict2Video2WorldAspectRatio.__args__,
        default="16:9",
        type=str,
        help="Aspect ratio of the generated output (width:height)",
    )
    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
    parser.add_argument(
        "--save_path",
        type=str,
        default="output/generated_image.jpg",
        help="Path to save the generated image (include file extension)",
    )
    parser.add_argument("--use_cuda_graphs", action="store_true", help="Use CUDA Graphs for the inference.")
    parser.add_argument("--disable_guardrail", action="store_true", help="Disable guardrail checks on prompts")
    parser.add_argument("--offload_guardrail", action="store_true", help="Offload guardrail to CPU to save GPU memory")
    parser.add_argument(
        "--benchmark",
        action="store_true",
        help="Run the generation in benchmark mode. It means that generation will be rerun a few times and the average generation time will be shown.",
    )
    parser.add_argument(
        "--use_fast_tokenizer",
        action="store_true",
        help="Use fast tokenizer for generation.",
    )

    # ---- Sparse Backend Integration ----
    add_pisa_args, add_svg_args, add_radial_args, add_sito_args, add_itm_args, add_fastercache_args, _ = _import_sparse_backends()
    add_pisa_args(parser)
    add_svg_args(parser)
    add_radial_args(parser)
    add_sito_args(parser)
    add_itm_args(parser)
    add_fastercache_args(parser)
    # -----------------------------------

    return parser.parse_args()


def setup_pipeline(args: argparse.Namespace, text_encoder: CosmosTextEncoder | None = None) -> Text2ImagePipeline:
    print_environment_info(args)
    if getattr(args, "use_fastercache", False) and getattr(args, "use_cuda_graphs", False):
        raise ValueError("[FasterCache] DreamGen FasterCache v1 is incompatible with --use_cuda_graphs.")

    config = get_cosmos_predict2_text2image_pipeline(
        model_size=args.model_size,
        fast_tokenizer=args.use_fast_tokenizer,
        pisa=getattr(args, "use_pisa", False),
        pisa_density=getattr(args, "pisa_density", 0.5),
        pisa_block_size=getattr(args, "pisa_block_size", 64),
        pisa_start_layer_idx=getattr(args, "pisa_start_layer_idx", 0),
        pisa_use_bias=getattr(args, "pisa_use_bias", False),
        radial=getattr(args, "use_radial", False),
        radial_decay_factor=getattr(args, "radial_decay_factor", 1.0),
        radial_block_size=getattr(args, "radial_block_size", 64),
        radial_start_layer_idx=getattr(args, "radial_start_layer_idx", 4),
        radial_model_type=getattr(args, "radial_model_type", "hunyuan"),
        itm=getattr(args, "use_itm", False),
        itm_compress_ratio=getattr(args, "itm_compress_ratio", None),
        itm_prune_from_step=getattr(args, "itm_prune_from_step", 1),
        itm_merge_from_step=getattr(args, "itm_merge_from_step", 2),
        itm_merge_attn=getattr(args, "itm_merge_attn", True),
        itm_merge_crossattn=getattr(args, "itm_merge_crossattn", False),
        itm_merge_mlp=getattr(args, "itm_merge_mlp", False),
        itm_start_layer_idx=getattr(args, "itm_start_layer_idx", None),
        itm_keep_last_n_dense=getattr(args, "itm_keep_last_n_dense", 2),
        sito=getattr(args, "use_sito", False),
        sito_prune_ratio=getattr(args, "sito_prune_ratio", None),
        sito_start_layer_idx=getattr(args, "sito_start_layer_idx", None),
        sito_keep_last_n_dense=getattr(args, "sito_keep_last_n_dense", 2),
        sito_patch_h=getattr(args, "sito_patch_h", 2),
        sito_patch_w=getattr(args, "sito_patch_w", 2),
        sito_noise_alpha=getattr(args, "sito_noise_alpha", 0.1),
        sito_sim_beta=getattr(args, "sito_sim_beta", 1.0),
        svg=getattr(args, "use_svg", False),
        svg_variant=getattr(args, "svg_variant", "svg1"),
        svg_start_layer_idx=getattr(args, "svg_start_layer_idx", 4),
        svg_dense_step_frac=getattr(args, "svg_dense_step_frac", 0.1),
        svg_sparsity=getattr(args, "svg_sparsity", 0.25),
        svg_num_sampled_rows=getattr(args, "svg_num_sampled_rows", 64),
        svg_sample_mse_max_row=getattr(args, "svg_sample_mse_max_row", 3000),
        svg_num_q_centroids=getattr(args, "svg_num_q_centroids", 50),
        svg_num_k_centroids=getattr(args, "svg_num_k_centroids", 200),
        svg_top_p_kmeans=getattr(args, "svg_top_p_kmeans", 0.9),
        svg_min_kc_ratio=getattr(args, "svg_min_kc_ratio", 0.0),
        svg_kmeans_iter_init=getattr(args, "svg_kmeans_iter_init", 5),
        svg_kmeans_iter_step=getattr(args, "svg_kmeans_iter_step", 2),
        fastercache=getattr(args, "use_fastercache", False),
        fastercache_start_step=getattr(args, "fastercache_start_step", 0),
        fastercache_model_interval=getattr(args, "fastercache_model_interval", 5),
        fastercache_block_interval=getattr(args, "fastercache_block_interval", 3),
    )
    if hasattr(args, "dit_path") and args.dit_path:
        dit_path = args.dit_path
    else:
        dit_path = get_cosmos_predict2_text2image_checkpoint(
            model_size=args.model_size, fast_tokenizer=args.use_fast_tokenizer, distilled=args.distill_steps > 0
        )
    log.info(f"Using dit_path: {dit_path}")

    # Disable guardrail if requested
    if args.disable_guardrail:
        log.warning("Guardrail checks are disabled")
        config.guardrail_config.enabled = False
    config.guardrail_config.offload_model_to_cpu = args.offload_guardrail

    misc.set_random_seed(seed=args.seed, by_rank=True)
    # Initialize cuDNN.
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True
    # Floating-point precision settings.
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cuda.matmul.allow_tf32 = True

    # Save config
    output_path = os.path.splitext(args.save_path)[0]
    output_dir = os.path.dirname(output_path)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
    LazyConfig.save_yaml(config, f"{output_path}.yaml")

    # Check if we're in a distributed environment (called from text2world)
    is_distributed = parallel_state.is_initialized() and torch.distributed.is_initialized()

    if is_distributed:
        # We're in a multi-GPU text2world context - only initialize on rank 0
        from imaginaire.utils.distributed import get_rank

        rank = get_rank()

        if rank == 0:
            log.info("Rank 0: Initializing Text2ImagePipeline for text2world")
            # Load models only on rank 0
            log.info(f"Initializing Text2ImagePipeline with model size: {args.model_size}")
            pipe = Text2ImagePipeline.from_config(
                config=config,
                dit_path=dit_path,
                device="cuda",
                torch_dtype=torch.bfloat16,
                load_ema_to_reg=args.load_ema,
                distill_steps=args.distill_steps,
            )
            fastercache_config = getattr(config.net, "fastercache_config", None)
            if fastercache_config is not None:
                *_, apply_fastercache = _import_sparse_backends()
                apply_fastercache(pipe.dit, fastercache_config, cfg_mode="sequential")

            return pipe
        else:
            log.info(f"Rank {rank}: Skipping Text2ImagePipeline initialization - will wait for rank 0")
            return None  # Return None for non-rank-0 processes
    else:
        # We're running as standalone text2image script
        # Only initialize distributed if num_gpus > 1 AND we're running standalone
        if hasattr(args, "num_gpus") and args.num_gpus > 1:
            log.info(f"Initializing distributed environment with {args.num_gpus} GPUs for context parallelism")

            # Check if distributed environment is already initialized
            if not parallel_state.is_initialized():
                distributed.init()
                parallel_state.initialize_model_parallel(context_parallel_size=args.num_gpus)
                log.info(f"Context parallel group initialized with {args.num_gpus} GPUs")
            else:
                log.info("Distributed environment already initialized, skipping initialization")
                # Check if we need to reinitialize with different context parallel size
                current_cp_size = parallel_state.get_context_parallel_world_size()
                if current_cp_size != args.num_gpus:
                    log.warning(f"Context parallel size mismatch: current={current_cp_size}, requested={args.num_gpus}")
                    log.warning("Using existing context parallel configuration")
                else:
                    log.info(f"Using existing context parallel group with {current_cp_size} GPUs")

        # Load models for standalone execution
        log.info(f"Initializing Text2ImagePipeline with model size: {args.model_size}")
        pipe = Text2ImagePipeline.from_config(
            config=config,
            dit_path=dit_path,
            use_text_encoder=text_encoder is None,
            device="cuda",
            torch_dtype=torch.bfloat16,
            load_ema_to_reg=args.load_ema,
            distill_steps=args.distill_steps,
        )

        # Set the provided text encoder if one was passed
        if text_encoder is not None:
            pipe.text_encoder = text_encoder

        fastercache_config = getattr(config.net, "fastercache_config", None)
        if fastercache_config is not None:
            *_, apply_fastercache = _import_sparse_backends()
            apply_fastercache(pipe.dit, fastercache_config, cfg_mode="sequential")

        return pipe


def process_single_generation(
    pipe: Text2ImagePipeline,
    prompt: str,
    output_path: str,
    negative_prompt: str,
    aspect_ratio: str,
    seed: int,
    use_cuda_graphs: bool,
    benchmark: bool,
) -> bool:
    log.info(f"Running Text2ImagePipeline\nprompt: {prompt}")

    # When benchmarking, run inference 4 times, exclude the 1st due to warmup and average time.
    num_repeats = 4 if benchmark else 1
    time_sum = 0
    for i in range(num_repeats):
        # Generate image
        if benchmark and i > 0:
            torch.cuda.synchronize()
            start_time = time.time()
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            aspect_ratio=aspect_ratio,
            seed=seed,
            use_cuda_graphs=use_cuda_graphs,
        )
        if benchmark and i > 0:
            torch.cuda.synchronize()
            elapsed = time.time() - start_time
            time_sum += elapsed
            log.info(f"[iter {i} / {num_repeats - 1}] Generation time: {elapsed:.1f} seconds.")
    if benchmark:
        time_avg = time_sum / (num_repeats - 1)
        log.critical(f"The benchmarked generation time for Text2ImagePipeline is {time_avg:.1f} seconds.")

    if image is not None:
        # save the generated image
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        log.info(f"Saving generated image to: {output_path}")
        save_image_or_video(image, output_path)
        log.success(f"Successfully saved image to: {output_path}")
        # save the prompts used to generate the video
        output_prompt_path = os.path.splitext(output_path)[0] + ".txt"
        prompts_to_save = {"prompt": prompt, "negative_prompt": negative_prompt}
        save_text_prompts(prompts_to_save, output_prompt_path)
        log.success(f"Successfully saved prompt file to: {output_prompt_path}")
        return True
    return False


def generate_image(args: argparse.Namespace, pipe: Text2ImagePipeline) -> None:
    if args.benchmark:
        log.warning(
            "Running in benchmark mode. Each generation will be rerun a couple of times and the average generation time will be shown."
        )
    # Text-to-image
    if args.batch_input_json is not None:
        # Process batch inputs from JSON file
        log.info(f"Loading batch inputs from JSON file: {args.batch_input_json}")
        with open(args.batch_input_json) as f:
            batch_inputs = json.load(f)

        for idx, item in enumerate(batch_inputs):
            log.info(f"Processing batch item {idx + 1}/{len(batch_inputs)}")
            prompt = item.get("prompt", "")
            output_image = item.get("output_image", f"output_{idx}.jpg")

            if not prompt:
                log.warning(f"Skipping item {idx}: Missing prompt")
                continue

            process_single_generation(
                pipe=pipe,
                prompt=prompt,
                output_path=output_image,
                negative_prompt=args.negative_prompt,
                aspect_ratio=args.aspect_ratio,
                seed=args.seed,
                use_cuda_graphs=args.use_cuda_graphs,
                benchmark=args.benchmark,
            )
    else:
        if args.use_cuda_graphs:
            log.warning(
                "Using CUDA Graphs for a single inference call may not be beneficial because of overhead of Graphs creation."
            )
        process_single_generation(
            pipe=pipe,
            prompt=args.prompt,
            output_path=args.save_path,
            negative_prompt=args.negative_prompt,
            aspect_ratio=args.aspect_ratio,
            seed=args.seed,
            use_cuda_graphs=args.use_cuda_graphs,
            benchmark=args.benchmark,
        )

    return


def cleanup_distributed():
    """Clean up the distributed environment if initialized."""
    if parallel_state.is_initialized():
        parallel_state.destroy_model_parallel()
        if torch.distributed.is_initialized():
            torch.distributed.destroy_process_group()


if __name__ == "__main__":
    args = parse_args()
    try:
        pipe = setup_pipeline(args)
        generate_image(args, pipe)
    finally:
        # Make sure to clean up the distributed environment
        cleanup_distributed()