# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""
Example script for image editing with OmniGen2.

    python image_edit.py \
        --image input.png \
        --model "OmniGen2/OmniGen2" \
        --prompt "Change the background to classroom." \
        --negative-prompt "(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar" \
        --num-inference-steps 50 \
        --seed 0 \
        --guidance-scale 5.0 \
        --guidance-scale-2 2.0 \
        --output outputs/image_edit.png \
        --num-outputs-per-prompt 2

    Note: For OmniGen2, `guidance_scale` works as `text_guidance_scale`,
    and `guidance_scale_2` works as `image_guidance_scale`.

Example script for image editing with Qwen-Image-Edit.

Usage (single image):
    python image_edit.py \
        --image input.png \
        --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
        --output output_image_edit.png \
        --num-inference-steps 50 \
        --cfg-scale 4.0 \
        --guidance-scale 1.0

Usage (multiple images):
    python image_edit.py \
        --image input1.png input2.png input3.png \
        --prompt "Combine these images into a single scene" \
        --output output_image_edit.png \
        --num-inference-steps 50 \
        --cfg-scale 4.0 \
        --guidance-scale 1.0

Usage (with cache-dit acceleration):
    python image_edit.py \
        --image input.png \
        --prompt "Edit description" \
        --cache-backend cache_dit \
        --cache-dit-max-continuous-cached-steps 3 \
        --cache-dit-residual-diff-threshold 0.24 \
        --cache-dit-enable-taylorseer

Usage (with tea_cache acceleration):
    python image_edit.py \
        --image input.png \
        --prompt "Edit description" \
        --cache-backend tea_cache \
        --tea-cache-rel-l1-thresh 0.25

Usage (layered):
    python image_edit.py \
        --model "Qwen/Qwen-Image-Layered" \
        --image input.png \
        --prompt "" \
        --output "layered" \
        --num-inference-steps 50 \
        --cfg-scale 4.0 \
        --layers 4 \
        --color-format "RGBA"

Usage (with CFG Parallel):
    python image_edit.py \
        --image input.png \
        --prompt "Edit description" \
        --cfg-parallel-size 2 \
        --num-inference-steps 50 \
        --cfg-scale 4.0

Usage (disable torch.compile):
    python image_edit.py \
        --image input.png \
        --prompt "Edit description" \
        --enforce-eager \
        --num-inference-steps 50 \
        --cfg-scale 4.0

For more options, run:
    python image_edit.py --help
"""

import argparse
import json
import os
import time
from pathlib import Path
from typing import Any

import torch
from PIL import Image

from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import (
    build_image_to_image_prompt,
    get_extra_body_params,
    get_model_class_name,
    should_init_extra_args_for_non_diffusion_stages,
)
from vllm_omni.platforms import current_omni_platform


def parse_profiler_config(value: str) -> dict[str, Any]:
    try:
        config = json.loads(value)
    except json.JSONDecodeError as e:
        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
    if not isinstance(config, dict):
        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
    return config


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
    parser.add_argument(
        "--model",
        default="Qwen/Qwen-Image-Edit",
        help=(
            "Diffusion model name or local path. "
            "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or Qwen/Qwen-Image-Edit-2511"
            "which supports QwenImageEditPlusPipeline."
        ),
    )
    parser.add_argument(
        "--image",
        type=str,
        nargs="+",
        required=True,
        help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.",
    )
    parser.add_argument(
        "--deploy-config",
        type=str,
        default=None,
        help=(
            "Path to a deploy YAML. Required for multi-stage image-edit pipelines "
            "whose deploy config is not auto-loaded."
        ),
    )
    parser.add_argument(
        "--prompt",
        type=str,
        required=True,
        help="Text prompt describing the edit to make to the image.",
    )
    parser.add_argument(
        "--negative-prompt",
        type=str,
        default=None,
        required=False,
    )
    parser.add_argument(
        "--width",
        type=int,
        default=None,
        metavar="W",
        help="Output image width in pixels. Default: None (pipeline's default).",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=None,
        metavar="H",
        help="Output image height in pixels. Default: None (pipeline's default).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=0,
        help="Random seed for deterministic results.",
    )
    parser.add_argument(
        "--cfg-scale",
        type=float,
        default=4.0,
        help=(
            "True classifier-free guidance scale (default: 4.0). Guidance scale as defined in Classifier-Free "
            "Diffusion Guidance. Classifier-free guidance is enabled by setting cfg_scale > 1 and providing "
            "a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, "
            "usually at the expense of lower image quality."
        ),
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=1.0,
        help=(
            "Guidance scale for guidance-distilled models (default: 1.0, disabled). "
            "Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale "
            "directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
        ),
    )
    parser.add_argument(
        "--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation."
    )
    parser.add_argument(
        "--extra-args",
        type=parse_profiler_config,
        default=None,
        help="JSON object copied to OmniDiffusionSamplingParams.extra_args, e.g. '{\"cfg_text_scale\": 4.0}'.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="output_image_edit.png",
        help=("Path to save the edited image (PNG). Or prefix for Qwen-Image-Layered model save images(PNG)."),
    )
    parser.add_argument(
        "--num-outputs-per-prompt",
        type=int,
        default=1,
        help="Number of images to generate for the given prompt.",
    )
    parser.add_argument(
        "--num-inference-steps",
        type=int,
        default=50,
        help="Number of denoising steps for the diffusion sampler.",
    )
    parser.add_argument(
        "--cache-backend",
        type=str,
        default=None,
        choices=["cache_dit", "tea_cache"],
        help=(
            "Cache backend to use for acceleration. "
            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
            "Default: None (no cache acceleration)."
        ),
    )
    parser.add_argument(
        "--ulysses-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ulysses sequence parallelism.",
    )
    parser.add_argument(
        "--ulysses-mode",
        type=str,
        default="strict",
        choices=["strict", "advanced_uaa"],
        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
    )
    parser.add_argument(
        "--ring-degree",
        type=int,
        default=1,
        help="Number of GPUs used for ring sequence parallelism.",
    )
    parser.add_argument(
        "--tensor-parallel-size",
        type=int,
        default=1,
        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
    )
    parser.add_argument(
        "--enable-expert-parallel",
        action="store_true",
        help="Enable expert parallelism for MoE layers.",
    )
    parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
    parser.add_argument(
        "--resolution",
        type=int,
        default=None,
        help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
    )

    parser.add_argument(
        "--color-format",
        type=str,
        default="RGB",
        help="For Qwen-Image-Layered, set to RGBA.",
    )

    # Cache-DiT specific parameters
    parser.add_argument(
        "--cache-dit-fn-compute-blocks",
        type=int,
        default=1,
        help="[cache-dit] Number of forward compute blocks. Optimized for single-transformer models.",
    )
    parser.add_argument(
        "--cache-dit-bn-compute-blocks",
        type=int,
        default=0,
        help="[cache-dit] Number of backward compute blocks.",
    )
    parser.add_argument(
        "--cache-dit-max-warmup-steps",
        type=int,
        default=4,
        help="[cache-dit] Maximum warmup steps (works for few-step models).",
    )
    parser.add_argument(
        "--cache-dit-residual-diff-threshold",
        type=float,
        default=0.24,
        help="[cache-dit] Residual diff threshold. Higher values enable more aggressive caching.",
    )
    parser.add_argument(
        "--cache-dit-max-continuous-cached-steps",
        type=int,
        default=3,
        help="[cache-dit] Maximum continuous cached steps to prevent precision degradation.",
    )
    parser.add_argument(
        "--cache-dit-enable-taylorseer",
        action="store_true",
        default=False,
        help="[cache-dit] Enable TaylorSeer acceleration (not suitable for few-step models).",
    )
    parser.add_argument(
        "--cache-dit-taylorseer-order",
        type=int,
        default=1,
        help="[cache-dit] TaylorSeer polynomial order.",
    )
    parser.add_argument(
        "--cache-dit-scm-steps-mask-policy",
        type=str,
        default=None,
        choices=[None, "slow", "medium", "fast", "ultra"],
        help="[cache-dit] SCM mask policy: None (disabled), slow, medium, fast, ultra.",
    )
    parser.add_argument(
        "--cache-dit-scm-steps-policy",
        type=str,
        default="dynamic",
        choices=["dynamic", "static"],
        help="[cache-dit] SCM steps policy: dynamic or static.",
    )

    # TeaCache specific parameters
    parser.add_argument(
        "--tea-cache-rel-l1-thresh",
        type=float,
        default=0.2,
        help="[tea_cache] Threshold for accumulated relative L1 distance.",
    )
    parser.add_argument(
        "--cfg-parallel-size",
        type=int,
        default=1,
        choices=[1, 2, 3],
        help="Number of GPUs used for classifier free guidance parallel size (max 3 branches).",
    )
    parser.add_argument(
        "--enforce-eager",
        action="store_true",
        default=None,
        help=(
            "Disable torch.compile and force eager execution. Left unset (None) "
            "so it is only forwarded when explicitly given; "
            "otherwise the per-stage deploy YAML value wins."
        ),
    )
    parser.add_argument(
        "--vae-use-slicing",
        action="store_true",
        help="Enable VAE slicing for memory optimization.",
    )
    parser.add_argument(
        "--vae-use-tiling",
        action="store_true",
        help="Enable VAE tiling for memory optimization.",
    )
    parser.add_argument(
        "--enable-cpu-offload",
        action="store_true",
        help="Enable CPU offloading for diffusion models.",
    )
    parser.add_argument(
        "--enable-layerwise-offload",
        action="store_true",
        help="Enable layerwise (blockwise) offloading on DiT modules.",
    )
    parser.add_argument(
        "--vae-patch-parallel-size",
        type=int,
        default=1,
        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
    )
    parser.add_argument(
        "--use-hsdp",
        action="store_true",
        help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
    )
    parser.add_argument(
        "--hsdp-shard-size",
        type=int,
        default=1,
        help="Number of GPUs to shard weights across for HSDP.",
    )
    parser.add_argument(
        "--hsdp-replicate-size",
        type=int,
        default=1,
        help="Number of HSDP replica groups.",
    )
    parser.add_argument(
        "--enable-diffusion-pipeline-profiler",
        action="store_true",
        help="Enable diffusion pipeline profiler to display stage durations.",
    )
    parser.add_argument(
        "--profiler-config",
        type=parse_profiler_config,
        default=None,
        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
    )
    return parser.parse_args()


def main():
    args = parse_args()

    if args.resolution and (args.width or args.height):
        raise ValueError("--resolution and --width/--height cannot be specified together")
    if args.width is not None and args.width <= 0:
        raise ValueError("--width must be a positive integer")
    if args.height is not None and args.height <= 0:
        raise ValueError("--height must be a positive integer")
    if not args.width and not args.height and not args.resolution:
        args.resolution = 640

    # Validate input images exist and load them
    input_images = []
    for image_path in args.image:
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Input image not found: {image_path}")

        img = Image.open(image_path).convert(args.color_format)
        input_images.append(img)

    # Use single image or list based on number of inputs
    if len(input_images) == 1:
        input_image = input_images[0]
    else:
        input_image = input_images

    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)

    parallel_config = DiffusionParallelConfig(
        ulysses_degree=args.ulysses_degree,
        ring_degree=args.ring_degree,
        cfg_parallel_size=args.cfg_parallel_size,
        tensor_parallel_size=args.tensor_parallel_size,
        enable_expert_parallel=args.enable_expert_parallel,
    )

    # Configure cache based on backend type
    cache_config = None
    if args.cache_backend == "cache_dit":
        # cache-dit configuration: Hybrid DBCache + SCM + TaylorSeer
        cache_config = {
            "Fn_compute_blocks": args.cache_dit_fn_compute_blocks,
            "Bn_compute_blocks": args.cache_dit_bn_compute_blocks,
            "max_warmup_steps": args.cache_dit_max_warmup_steps,
            "residual_diff_threshold": args.cache_dit_residual_diff_threshold,
            "max_continuous_cached_steps": args.cache_dit_max_continuous_cached_steps,
            "enable_taylorseer": args.cache_dit_enable_taylorseer,
            "taylorseer_order": args.cache_dit_taylorseer_order,
            "scm_steps_mask_policy": args.cache_dit_scm_steps_mask_policy,
            "scm_steps_policy": args.cache_dit_scm_steps_policy,
        }
    elif args.cache_backend == "tea_cache":
        # TeaCache configuration
        cache_config = {
            "rel_l1_thresh": args.tea_cache_rel_l1_thresh,
            # Note: coefficients will use model-specific defaults based on model_type
        }

    # Initialize Omni with appropriate pipeline
    omni_kwargs: dict[str, Any] = dict(
        model=args.model,
        enable_layerwise_offload=args.enable_layerwise_offload,
        vae_use_slicing=args.vae_use_slicing,
        vae_use_tiling=args.vae_use_tiling,
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        parallel_config=parallel_config,
        enable_cpu_offload=args.enable_cpu_offload,
        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
        profiler_config=args.profiler_config,
    )
    if args.enforce_eager is not None:
        omni_kwargs["enforce_eager"] = args.enforce_eager
    if args.deploy_config:
        omni_kwargs["deploy_config"] = args.deploy_config
    omni = Omni(**omni_kwargs)
    model_class_name = get_model_class_name(omni)
    declared_extra_body_params = get_extra_body_params(model_class_name)
    print("Pipeline loaded")

    profiler_enabled = args.profiler_config is not None

    # Time profiling for generation
    print(f"\n{'=' * 60}")
    print("Generation Configuration:")
    print(f"  Model: {args.model}")
    print(f"  Inference steps: {args.num_inference_steps}")
    print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
    if args.height is not None or args.width is not None:
        print(f"  Output size: {args.width or 'auto'}x{args.height or 'auto'}")
    if isinstance(input_image, list):
        print(f"  Number of input images: {len(input_image)}")
        for idx, img in enumerate(input_image):
            print(f"    Image {idx + 1} size: {img.size}")
    else:
        print(f"  Input image size: {input_image.size}")
    print(
        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
    )
    print(f"{'=' * 60}\n")

    generation_start = time.perf_counter()

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
        omni.start_profile()

    prompt_dict = build_image_to_image_prompt(
        model_class_name=model_class_name,
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        input_image=input_image,
        height=args.height,
        width=args.width,
    )

    extra_args_from_cli = dict(args.extra_args or {})
    if args.negative_prompt is not None:
        extra_args_from_cli.setdefault("negative_prompt", args.negative_prompt)

    diffusion_params = OmniDiffusionSamplingParams(
        generator=generator,
        true_cfg_scale=args.cfg_scale,
        guidance_scale=args.guidance_scale,
        guidance_scale_2=args.guidance_scale_2,
        num_inference_steps=args.num_inference_steps,
        num_outputs_per_prompt=args.num_outputs_per_prompt,
        layers=args.layers,
        resolution=args.resolution,
        height=args.height,
        width=args.width,
    )
    if declared_extra_body_params:
        apply_declared_extra_args(
            diffusion_params,
            declared_extra_body_params,
            extra_args_from_cli,
        )
    else:
        diffusion_params.extra_args.update({k: v for k, v in extra_args_from_cli.items() if v is not None})

    # Build per-stage sampling params for multi-stage models
    init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
        model_class_name,
    )
    defaults = list(omni.default_sampling_params_list or [])
    sampling_params_list = [clone_sampling_params(p) for p in defaults]
    if not sampling_params_list:
        sampling_params_list = [diffusion_params]

    diffusion_replaced = False
    for idx, params in enumerate(sampling_params_list):
        if isinstance(params, OmniDiffusionSamplingParams):
            merged_extra = dict(getattr(params, "extra_args", {}) or {})
            merged_extra.update(diffusion_params.extra_args)
            diffusion_params.extra_args = merged_extra
            sampling_params_list[idx] = diffusion_params
            diffusion_replaced = True
        elif init_non_diffusion and hasattr(params, "extra_args"):
            if params.extra_args is None:
                params.extra_args = {}

    if not diffusion_replaced and len(sampling_params_list) == 1:
        sampling_params_list = [diffusion_params]

    outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start

    # Print profiling results
    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")

    if profiler_enabled:
        print("\n[Profiler] Stopping profiler and collecting results...")
        profile_results = omni.stop_profile()
        if profile_results and isinstance(profile_results, dict):
            traces = profile_results.get("traces", [])
            print("\n" + "=" * 60)
            print("PROFILING RESULTS:")
            for rank, trace in enumerate(traces):
                print(f"\nRank {rank}:")
                if trace:
                    print(f"  • Trace: {trace}")
            if not traces:
                print("  No traces collected.")
            print("=" * 60)
        else:
            print("[Profiler] No valid profiling data returned.")

    if not outputs:
        raise ValueError("No output generated from omni.generate()")

    images = None
    for output in outputs:
        images = getattr(output, "images", None)
        if images:
            break
        req_out = getattr(output, "request_output", None)
        images = getattr(req_out, "images", None) if req_out is not None else None
        if images:
            break

    if not images:
        raise ValueError("No images found in request_output")

    # Save output image(s)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    suffix = output_path.suffix or ".png"
    stem = output_path.stem or "output_image_edit"

    # Handle layered output (each image may be a list of layers)
    if args.num_outputs_per_prompt <= 1:
        img = images[0]
        # Check if this is a layered output (list of images)
        if isinstance(img, list):
            for sub_idx, sub_img in enumerate(img):
                save_path = output_path.parent / f"{stem}_{sub_idx}{suffix}"
                sub_img.save(save_path)
                print(f"Saved edited image to {os.path.abspath(save_path)}")
        else:
            img.save(output_path)
            print(f"Saved edited image to {os.path.abspath(output_path)}")
    else:
        for idx, img in enumerate(images):
            # Check if this is a layered output (list of images)
            if isinstance(img, list):
                for sub_idx, sub_img in enumerate(img):
                    save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}"
                    sub_img.save(save_path)
                    print(f"Saved edited image to {os.path.abspath(save_path)}")
            else:
                save_path = output_path.parent / f"{stem}_{idx}{suffix}"
                img.save(save_path)
                print(f"Saved edited image to {os.path.abspath(save_path)}")


if __name__ == "__main__":
    main()