vllm-omni/tests/diffusion/models/joy_image/test_joy_image.py

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

import json
import logging
from types import SimpleNamespace

import pytest
import torch
import torch.nn.functional as F
from PIL import Image

from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
from vllm_omni.diffusion.models.joy_image import (
    pipeline_joy_image_edit as joy_pipeline_module,
)
from vllm_omni.diffusion.models.joy_image.cfg_parallel import JoyImageEditCFGParallelMixin
from vllm_omni.diffusion.models.joy_image.joy_image_edit_transformer import (
    JoyImageAttention,
    JoyImageEditTransformer3DModel,
    _raise_if_unsupported_tensor_parallel,
)
from vllm_omni.diffusion.models.joy_image.pipeline_joy_image_edit import (
    JOY_MAX_IMAGE_SEQ_LEN,
    JOY_PROMPT_TEMPLATE,
    JOY_VISION_TOKEN,
    JoyImageEditPipeline,
    _cast_floating_model_inputs,
    _format_qwen_multimodal_prompt,
    _get_transformer_config_kwargs_from_od_config,
    _raise_if_unsupported_hsdp,
    _resize_center_crop,
    _should_defer_component_device_placement,
    _uses_model_level_cpu_offload,
    get_joy_image_edit_pre_process_func,
    retrieve_timesteps,
)
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


def _torch_sdpa_diffusion_config():
    return SimpleNamespace(
        diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")),
        parallel_config=SimpleNamespace(ring_degree=1),
        diffusion_kv_cache_dtype=None,
        diffusion_kv_cache_skip_step_indices=None,
        diffusion_kv_cache_skip_layer_indices=None,
    )


def _make_joy_attention(*, dtype: torch.dtype = torch.float32) -> JoyImageAttention:
    with set_current_diffusion_config(_torch_sdpa_diffusion_config()):
        attention = JoyImageAttention(
            dim=32,
            num_attention_heads=4,
            attention_head_dim=8,
            prefix="double_blocks.0.attn",
        )
    return attention.to(dtype=dtype)


def _reference_joy_attention_output(
    attention: JoyImageAttention,
    hidden_states: torch.Tensor,
    encoder_hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    img_query, img_key, img_value = attention._stream_qkv(
        attention.img_attn_qkv,
        attention.img_attn_q_norm,
        attention.img_attn_k_norm,
        hidden_states,
    )
    txt_query, txt_key, txt_value = attention._stream_qkv(
        attention.txt_attn_qkv,
        attention.txt_attn_q_norm,
        attention.txt_attn_k_norm,
        encoder_hidden_states,
    )
    joint_query = torch.cat([img_query, txt_query], dim=1)
    joint_key = torch.cat([img_key, txt_key], dim=1)
    joint_value = torch.cat([img_value, txt_value], dim=1)

    if attention_mask is not None and attention_mask.ndim == 2:
        attention_mask = attention_mask.to(torch.bool)[:, None, None, :]

    joint_hidden_states = F.scaled_dot_product_attention(
        joint_query.transpose(1, 2),
        joint_key.transpose(1, 2),
        joint_value.transpose(1, 2),
        attn_mask=attention_mask,
        dropout_p=0.0,
        is_causal=False,
        scale=1.0 / (attention.head_dim**0.5),
    )
    joint_hidden_states = joint_hidden_states.transpose(1, 2).flatten(2, 3).to(joint_query.dtype)

    image_seq_len = hidden_states.shape[1]
    return (
        attention.img_attn_proj(joint_hidden_states[:, :image_seq_len]),
        attention.txt_attn_proj(joint_hidden_states[:, image_seq_len:]),
    )


def _assert_attention_outputs_close(
    actual: tuple[torch.Tensor, torch.Tensor],
    expected: tuple[torch.Tensor, torch.Tensor],
    *,
    dtype: torch.dtype,
) -> None:
    atol = rtol = 5e-3 if dtype is torch.bfloat16 else 1e-5
    max_abs = max(
        (actual_item - expected_item).abs().max().item() for actual_item, expected_item in zip(actual, expected)
    )
    assert max_abs <= atol
    for actual_item, expected_item in zip(actual, expected):
        torch.testing.assert_close(actual_item, expected_item, atol=atol, rtol=rtol)


def _write_model_configs(tmp_path):
    (tmp_path / "vae").mkdir()
    (tmp_path / "transformer").mkdir()
    (tmp_path / "vae" / "config.json").write_text(
        json.dumps(
            {
                "scale_factor_spatial": 8,
                "z_dim": 16,
                "latents_mean": [0.0] * 16,
                "latents_std": [1.0] * 16,
            }
        )
    )
    (tmp_path / "transformer" / "config.json").write_text(
        json.dumps(
            {
                "patch_size": [1, 2, 2],
                "hidden_size": 32,
                "num_attention_heads": 4,
                "num_layers": 1,
                "in_channels": 4,
                "out_channels": 4,
                "text_dim": 16,
            }
        )
    )
    return SimpleNamespace(
        model=str(tmp_path),
        model_class_name="JoyImageEditPipeline",
    )


def _make_params(**overrides):
    guidance_scale = overrides.pop("guidance_scale", None)
    values = {
        "height": None,
        "width": None,
        "guidance_scale": guidance_scale if guidance_scale is not None else 1.0,
        "guidance_scale_provided": guidance_scale is not None,
        "true_cfg_scale": None,
        "num_outputs_per_prompt": 1,
        "num_inference_steps": None,
        "timesteps": None,
        "sigmas": None,
        "max_sequence_length": None,
        "latents": None,
        "output_type": None,
        "generator": None,
        "seed": None,
        "cfg_normalize": False,
        "guidance_scale_2": None,
        "do_classifier_free_guidance": False,
    }
    values.update(overrides)
    return SimpleNamespace(**values)


def _make_request(*, prompt=None, params=None):
    prompt = prompt or {
        "prompt": "make it brighter",
        "multi_modal_data": {"image": Image.new("RGB", (2048, 1024), color="white")},
    }
    params = params or _make_params()
    return SimpleNamespace(
        prompts=[prompt],
        sampling_params=params,
        request_id="joy-test",
        is_dummy_run=lambda: False,
    )


def _make_single_omni_request(*, prompt=None, params=None):
    prompt = prompt or {
        "prompt": "make it brighter",
        "multi_modal_data": {"image": Image.new("RGB", (2048, 1024), color="white")},
    }
    params = params or _make_params()
    return OmniDiffusionRequest(
        prompt=prompt,
        sampling_params=params,
        request_id="joy-test",
    )


class _TinyJoyLayerwisePipeline(torch.nn.Module):
    _dit_modules = ["transformer"]
    _encoder_modules: list[str] = []
    _vae_modules: list[str] = []
    _resident_modules: list[str] = []
    _offload_transformer_if_model_level = JoyImageEditPipeline._offload_transformer_if_model_level

    def __init__(self, *, device: torch.device, cleanup_mode: str):
        super().__init__()
        self.device = device
        self.scheduler = SimpleNamespace(
            set_timesteps=lambda num_inference_steps, device: setattr(
                self.scheduler,
                "timesteps",
                torch.tensor([1.0], device=device),
            )
        )
        self.latent_channels = 4
        self.od_config = SimpleNamespace(
            enable_cpu_offload=False,
            enable_layerwise_offload=True,
            pin_cpu_memory=False,
            parallel_config=SimpleNamespace(use_hsdp=False),
        )
        self.transformer = JoyImageEditTransformer3DModel(
            in_channels=4,
            out_channels=4,
            hidden_size=32,
            text_dim=16,
            num_layers=2,
            num_attention_heads=4,
            patch_size=(1, 2, 2),
        ).to(device=device, dtype=torch.bfloat16)
        self.vae = torch.nn.Linear(1, 1)
        self.cleanup_mode = cleanup_mode

    def encode_prompt(self, *args, **kwargs):
        return (
            torch.zeros(1, 5, 16, device=self.device, dtype=torch.bfloat16),
            torch.ones(1, 5, dtype=torch.long, device=self.device),
        )

    def resolve_effective_true_cfg_scale(self, req, default_true_cfg_scale=4.0):
        return JoyImageEditPipeline.resolve_effective_true_cfg_scale(
            req,
            default_true_cfg_scale=default_true_cfg_scale,
        )

    def _prepare_latents(self, **kwargs):
        latents = torch.randn(1, 2, 4, 1, 4, 4, device=self.device, dtype=torch.bfloat16)
        return latents, latents[:, :1].clone()

    def diffuse(self, **kwargs):
        return self.transformer(
            hidden_states=kwargs["latents"],
            timestep=kwargs["timesteps"][0],
            encoder_hidden_states=kwargs["prompt_embeds"],
            encoder_hidden_states_mask=kwargs["prompt_embeds_mask"],
            return_dict=False,
        )[0]

    def _decode_latents(self, latents):
        JoyImageEditPipeline._prepare_vae_for_decode(self)
        return latents[:, -1, :, 0]


def _make_layerwise_forward_request(*, dummy: bool = False):
    request = _make_request(
        prompt={
            "prompt": "make it brighter",
            "additional_information": {
                "image_tensor": torch.zeros(1, 3, 1, 16, 16),
                "prompt_image": Image.new("RGB", (16, 16)),
                "height": 16,
                "width": 16,
            },
        },
        params=_make_params(guidance_scale=1.0, num_inference_steps=1),
    )
    request.is_dummy_run = lambda: dummy
    return request


def _run_tiny_joy_layerwise_sequence(first_mode: str):
    device = torch.device("cuda:0")
    pipeline = _TinyJoyLayerwisePipeline(device=device, cleanup_mode=first_mode)
    backend = LayerWiseOffloadBackend(
        OffloadConfig(strategy=OffloadStrategy.LAYER_WISE, pin_cpu_memory=False),
        device,
    )
    backend.enable(pipeline)
    try:
        if first_mode == "dummy":
            first_request = _make_layerwise_forward_request(dummy=True)
            first_output_type = "pil"
        elif first_mode == "latent":
            first_request = _make_layerwise_forward_request(dummy=False)
            first_output_type = "latent"
        elif first_mode == "decoded":
            first_request = _make_layerwise_forward_request(dummy=False)
            first_output_type = "pil"
        else:
            raise AssertionError(f"unexpected first mode {first_mode!r}")

        with torch.inference_mode():
            JoyImageEditPipeline.forward(pipeline, first_request, output_type=first_output_type)

        # These non-block modules execute before double_blocks[0], so layerwise
        # hooks cannot repair them lazily after request cleanup.
        assert next(pipeline.transformer.img_in.parameters()).device == device
        assert next(pipeline.transformer.condition_embedder.parameters()).device == device
        assert next(pipeline.transformer.proj_out.parameters()).device == device

        second_request = _make_layerwise_forward_request(dummy=False)
        with torch.inference_mode():
            output = JoyImageEditPipeline.forward(pipeline, second_request, output_type="pil")
    finally:
        backend.disable()
        torch.accelerator.empty_cache()

    assert output.output.shape == (1, 4, 4, 4)


def test_joy_registry_and_metadata_entries(tmp_path):
    from vllm_omni.diffusion.registry import (
        DiffusionModelRegistry,
        get_diffusion_post_process_func,
        get_diffusion_pre_process_func,
    )

    od_config = _write_model_configs(tmp_path)
    model_class = DiffusionModelRegistry._try_load_model_cls("JoyImageEditPipeline")
    assert model_class is JoyImageEditPipeline
    assert callable(get_diffusion_pre_process_func(od_config))
    postprocess = get_diffusion_post_process_func(od_config)
    assert callable(postprocess)
    assert len(postprocess(torch.zeros(1, 3, 2, 2))) == 1

    metadata = get_diffusion_model_metadata("JoyImageEditPipeline")
    assert metadata.supports_multimodal_inputs is True
    assert metadata.max_multimodal_image_inputs == 1


def test_get_model_path_downloads_only_runtime_configs(monkeypatch):
    captured = {}

    def fake_download(model_name, revision, allow_patterns, **kwargs):
        captured.update(
            {
                "model_name": model_name,
                "revision": revision,
                "allow_patterns": allow_patterns,
                "kwargs": kwargs,
            }
        )
        return "/tmp/joy-configs"

    monkeypatch.setattr(
        joy_pipeline_module,
        "download_weights_from_hf_specific",
        fake_download,
    )

    assert joy_pipeline_module._get_model_path("repo/joy") == "/tmp/joy-configs"
    assert captured == {
        "model_name": "repo/joy",
        "revision": None,
        "allow_patterns": ["vae/config.json", "transformer/config.json"],
        "kwargs": {"require_all": True},
    }


def test_transformer_config_kwargs_from_od_config_filters_metadata():
    class FakeTransformerConfig:
        def to_dict(self):
            return {
                "_class_name": "JoyImageEditTransformer3DModel",
                "_diffusers_version": "0.38.0",
                "hidden_size": 32,
                "num_layers": 1,
                "num_attention_heads": 4,
            }

    od_config = SimpleNamespace(tf_model_config=FakeTransformerConfig())

    assert _get_transformer_config_kwargs_from_od_config(od_config) == {
        "hidden_size": 32,
        "num_layers": 1,
        "num_attention_heads": 4,
    }


def test_defer_component_device_placement_for_offload_only():
    assert (
        _should_defer_component_device_placement(
            SimpleNamespace(
                enable_cpu_offload=False,
                enable_layerwise_offload=False,
                parallel_config=SimpleNamespace(use_hsdp=False),
            )
        )
        is False
    )
    assert _should_defer_component_device_placement(SimpleNamespace(enable_cpu_offload=True)) is True
    assert _should_defer_component_device_placement(SimpleNamespace(enable_layerwise_offload=True)) is True
    assert (
        _should_defer_component_device_placement(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True)))
        is False
    )


@pytest.mark.parametrize(
    ("enable_cpu_offload", "enable_layerwise_offload", "expected"),
    [
        (True, False, True),
        (False, True, False),
        (True, True, False),
        (False, False, False),
    ],
)
def test_model_level_cpu_offload_predicate_matches_effective_strategy(
    enable_cpu_offload,
    enable_layerwise_offload,
    expected,
):
    od_config = SimpleNamespace(
        enable_cpu_offload=enable_cpu_offload,
        enable_layerwise_offload=enable_layerwise_offload,
    )

    assert _uses_model_level_cpu_offload(od_config) is expected


def test_model_level_cpu_offload_cleanup_still_moves_transformer(monkeypatch):
    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=True,
        enable_layerwise_offload=False,
        pin_cpu_memory=False,
    )
    pipeline.transformer = torch.nn.Linear(1, 1)
    move_calls = []

    def fake_move_params(module, target_device, **kwargs):
        move_calls.append((module, target_device, kwargs))

    monkeypatch.setattr(
        joy_pipeline_module.SequentialOffloadHook,
        "_move_params",
        fake_move_params,
    )
    monkeypatch.setattr(joy_pipeline_module.current_omni_platform, "empty_cache", lambda: None)

    JoyImageEditPipeline._offload_transformer_if_model_level(pipeline)

    assert move_calls == [
        (
            pipeline.transformer,
            torch.device("cpu"),
            {"non_blocking": False, "pin_memory": False},
        )
    ]


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@pytest.mark.parametrize("first_mode", ["dummy", "decoded", "latent"])
def test_layerwise_joy_lifecycle_survives_second_forward(first_mode):
    _run_tiny_joy_layerwise_sequence(first_mode)


def test_joy_hsdp_is_explicitly_unsupported():
    with pytest.raises(ValueError, match="does not support HSDP"):
        _raise_if_unsupported_hsdp(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True)))

    _raise_if_unsupported_hsdp(SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False)))
    _raise_if_unsupported_hsdp(SimpleNamespace())


def test_joy_tensor_parallelism_is_explicitly_unsupported():
    with pytest.raises(ValueError, match="do not support Tensor Parallelism"):
        _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2)))

    _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1)))
    _raise_if_unsupported_tensor_parallel(SimpleNamespace(parallel_config=SimpleNamespace()))
    _raise_if_unsupported_tensor_parallel(SimpleNamespace())
    _raise_if_unsupported_tensor_parallel(None)


def test_transformer_rejects_tensor_parallelism_before_building_layers():
    with pytest.raises(ValueError, match="do not support Tensor Parallelism"):
        JoyImageEditTransformer3DModel(
            od_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2)),
            in_channels=4,
            out_channels=4,
            hidden_size=32,
            text_dim=16,
            num_layers=1,
            num_attention_heads=4,
            patch_size=(1, 2, 2),
        )


def test_pipeline_rejects_tensor_parallelism_before_loading_components():
    with pytest.raises(ValueError, match="do not support Tensor Parallelism"):
        JoyImageEditPipeline(
            od_config=SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=2))
        )


def test_component_discovery_treats_vae_as_offload_peer():
    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.transformer = torch.nn.Linear(1, 1)
    pipeline.text_encoder = torch.nn.Linear(1, 1)
    pipeline.vae = torch.nn.Linear(1, 1)

    modules = ModuleDiscovery.discover(pipeline)

    assert modules.dit_names == ["transformer"]
    assert modules.encoder_names == ["text_encoder", "vae"]
    assert modules.vaes == []


def test_preprocess_requires_exactly_one_image(tmp_path):
    od_config = _write_model_configs(tmp_path)
    preprocess = get_joy_image_edit_pre_process_func(od_config)
    no_image = _make_request(prompt={"prompt": "x", "multi_modal_data": {}})
    with pytest.raises(ValueError, match="exactly one input image"):
        preprocess(no_image)

    multi_image = _make_request(
        prompt={
            "prompt": "x",
            "multi_modal_data": {"image": [Image.new("RGB", (32, 32)), Image.new("RGB", (32, 32))]},
        }
    )
    with pytest.raises(ValueError, match="exactly one image"):
        preprocess(multi_image)


def test_preprocess_rejects_batched_prompts_and_partial_size(tmp_path):
    od_config = _write_model_configs(tmp_path)
    preprocess = get_joy_image_edit_pre_process_func(od_config)
    batched = _make_request()
    batched.prompts.append(batched.prompts[0])
    with pytest.raises(ValueError, match="exactly one prompt"):
        preprocess(batched)

    partial_size = _make_request(params=_make_params(height=512, width=None))
    with pytest.raises(ValueError, match="both `height` and `width`"):
        preprocess(partial_size)


def test_preprocess_uses_diffusers_bucket_for_default_image_size(tmp_path):
    od_config = _write_model_configs(tmp_path)
    preprocess = get_joy_image_edit_pre_process_func(od_config)

    request = preprocess(_make_request())
    prompt = request.prompts[0]
    info = prompt["additional_information"]
    height = info["height"]
    width = info["width"]

    assert (height, width) == (704, 1408)
    assert (height // 16) * (width // 16) <= JOY_MAX_IMAGE_SEQ_LEN
    assert info["image_tensor"].shape == (1, 3, 1, height, width)
    assert info["original_size"] == (2048, 1024)
    assert info["resized_size"] == (width, height)
    assert request.sampling_params.height == height
    assert request.sampling_params.width == width


def test_preprocess_accepts_single_omni_request_prompt(tmp_path):
    od_config = _write_model_configs(tmp_path)
    preprocess = get_joy_image_edit_pre_process_func(od_config)

    request = preprocess(_make_single_omni_request())
    prompt = request.prompt
    info = prompt["additional_information"]

    assert isinstance(prompt, dict)
    assert not hasattr(request, "prompts")
    assert info["image_tensor"].shape == (1, 3, 1, 704, 1408)
    assert request.sampling_params.height == 704
    assert request.sampling_params.width == 1408


def test_resize_center_crop_crops_instead_of_stretching():
    image = Image.new("RGB", (4, 2))
    pixels = image.load()
    colors = [
        (255, 0, 0),
        (0, 255, 0),
        (0, 0, 255),
        (255, 255, 255),
    ]
    for x, color in enumerate(colors):
        for y in range(2):
            pixels[x, y] = color

    cropped = _resize_center_crop(image, height=2, width=2)
    cropped_pixels = cropped.load()

    assert cropped.size == (2, 2)
    assert [[cropped_pixels[x, y] for x in range(2)] for y in range(2)] == [
        [(0, 255, 0), (0, 0, 255)],
        [(0, 255, 0), (0, 0, 255)],
    ]


def test_joy_cfg_normalize_uses_channel_dimension():
    mixin = JoyImageEditCFGParallelMixin()
    noise_pred = torch.tensor([[[[[[3.0, 4.0]]], [[[4.0, 3.0]]]]]])
    comb_pred = torch.tensor([[[[[[6.0, 8.0]]], [[[8.0, 6.0]]]]]])

    normalized = mixin.cfg_normalize_function(noise_pred, comb_pred)
    expected = comb_pred * (
        torch.norm(noise_pred, dim=2, keepdim=True) / torch.norm(comb_pred, dim=2, keepdim=True).clamp_min(1e-6)
    )

    assert torch.equal(normalized, expected)


def test_joy_diffuse_preserves_scheduler_timestep_dtype_for_time_embedding():
    class FakeProgressBar:
        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            return False

        def update(self):
            pass

    class FakeScheduler:
        def set_begin_index(self, index):
            self.begin_index = index

    class FakeMixin(JoyImageEditCFGParallelMixin):
        def __init__(self):
            self.scheduler = FakeScheduler()
            self.captured_timestep_dtype = None
            self._interrupt = False

        @property
        def interrupt(self):
            return self._interrupt

        def progress_bar(self, total):
            return FakeProgressBar()

        def predict_noise_maybe_with_cfg(self, **kwargs):
            self.captured_timestep_dtype = kwargs["positive_kwargs"]["timestep"].dtype
            return torch.zeros_like(kwargs["positive_kwargs"]["hidden_states"])

        def scheduler_step_maybe_with_cfg(self, noise_pred, timestep, latents, do_true_cfg):
            return latents

    mixin = FakeMixin()
    latents = torch.zeros(1, 2, 4, 1, 4, 4, dtype=torch.bfloat16)
    image_latents = torch.zeros(1, 1, 4, 1, 4, 4, dtype=torch.bfloat16)

    mixin.diffuse(
        latents=latents,
        image_latents=image_latents,
        prompt_embeds=torch.zeros(1, 2, 4, dtype=torch.bfloat16),
        prompt_embeds_mask=torch.ones(1, 2, dtype=torch.long),
        negative_prompt_embeds=None,
        negative_prompt_embeds_mask=None,
        timesteps=torch.tensor([999.125], dtype=torch.float32),
        do_true_cfg=False,
        true_cfg_scale=1.0,
    )

    assert mixin.captured_timestep_dtype == torch.float32


def test_joy_diffuse_replaces_reference_latents_out_of_place():
    class FakeProgressBar:
        def __enter__(self):
            return self

        def __exit__(self, exc_type, exc_value, traceback):
            return False

        def update(self):
            pass

    class FakeScheduler:
        def set_begin_index(self, index):
            self.begin_index = index

    class FakeMixin(JoyImageEditCFGParallelMixin):
        def __init__(self):
            self.scheduler = FakeScheduler()
            self.scheduler_returned_latents = None
            self._interrupt = False

        @property
        def interrupt(self):
            return self._interrupt

        def progress_bar(self, total):
            return FakeProgressBar()

        def predict_noise_maybe_with_cfg(self, **kwargs):
            return torch.zeros_like(kwargs["positive_kwargs"]["hidden_states"])

        def scheduler_step_maybe_with_cfg(self, noise_pred, timestep, latents, do_true_cfg):
            stepped = latents.clone()
            stepped[:, :1] = -5
            self.scheduler_returned_latents = stepped
            return stepped

    mixin = FakeMixin()
    latents = torch.arange(3, dtype=torch.float32).view(1, 3, 1, 1, 1, 1)
    original_latents = latents.clone()
    image_latents = torch.full((1, 1, 1, 1, 1, 1), 100.0)

    output = mixin.diffuse(
        latents=latents,
        image_latents=image_latents,
        prompt_embeds=torch.zeros(1, 2, 4),
        prompt_embeds_mask=torch.ones(1, 2, dtype=torch.long),
        negative_prompt_embeds=None,
        negative_prompt_embeds_mask=None,
        timesteps=torch.tensor([999.125], dtype=torch.float32),
        do_true_cfg=False,
        true_cfg_scale=1.0,
    )

    assert torch.equal(latents, original_latents)
    assert mixin.scheduler_returned_latents is not None
    assert torch.equal(mixin.scheduler_returned_latents[:, :1], torch.full_like(image_latents, -5))
    assert torch.equal(output[:, :1], image_latents)


def test_format_qwen_multimodal_prompt_matches_diffusers_image_placeholder_replacement():
    formatted = _format_qwen_multimodal_prompt("make it brighter")

    assert formatted == (f"<|im_start|>user\n{JOY_VISION_TOKEN}make it brighter<|im_end|>\n")


def test_preprocess_maps_explicit_size_to_nearest_diffusers_bucket(tmp_path):
    od_config = _write_model_configs(tmp_path)
    preprocess = get_joy_image_edit_pre_process_func(od_config)
    params = _make_params(height=2048, width=2048)

    request = preprocess(_make_request(params=params))

    assert request.sampling_params.height == 1024
    assert request.sampling_params.width == 1024


@pytest.mark.parametrize(
    "params_kwargs, expected_scale",
    [
        ({}, 4.0),
        ({"guidance_scale": 3.5}, 3.5),
        ({"true_cfg_scale": 4.0}, 4.0),
        ({"guidance_scale": 4.0, "true_cfg_scale": 4.0}, 4.0),
    ],
)
def test_resolve_effective_true_cfg_scale(params_kwargs, expected_scale):
    request = _make_request(params=_make_params(**params_kwargs))
    assert JoyImageEditPipeline.resolve_effective_true_cfg_scale(request) == expected_scale


@pytest.mark.parametrize("guidance_scale", [1.0, 3.0])
def test_explicit_guidance_scale_conflicts_with_true_cfg_scale(guidance_scale):
    request = _make_request(params=_make_params(guidance_scale=guidance_scale, true_cfg_scale=4.0))
    with pytest.raises(ValueError, match="compatibility alias"):
        JoyImageEditPipeline.resolve_effective_true_cfg_scale(request)


def test_pad_prompt_embeds_keeps_last_tokens_and_builds_mask():
    first = torch.arange(20, dtype=torch.float32).reshape(5, 4)
    second = torch.arange(8, dtype=torch.float32).reshape(2, 4)

    prompt_embeds, prompt_mask = JoyImageEditPipeline._pad_prompt_embeds(
        [first, second],
        max_sequence_length=3,
    )

    assert prompt_embeds.shape == (2, 3, 4)
    assert torch.equal(prompt_embeds[0], first[-3:])
    assert torch.equal(prompt_embeds[1, :2], second)
    assert torch.equal(prompt_embeds[1, 2], torch.zeros(4))
    assert torch.equal(
        prompt_mask,
        torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long),
    )

    with pytest.raises(ValueError, match="greater than 0"):
        JoyImageEditPipeline._pad_prompt_embeds([first], max_sequence_length=0)


@pytest.mark.parametrize(
    ("prompt_embeds", "prompt_embeds_mask", "match"),
    [
        (
            torch.zeros(2, 4),
            torch.ones(2, 4, dtype=torch.long),
            "prompt_embeds must be a 3D tensor",
        ),
        (
            torch.zeros(2, 4, 8),
            torch.ones(2, 4, 1, dtype=torch.long),
            "prompt_embeds_mask must be a 2D tensor",
        ),
        (
            torch.zeros(2, 4, 8),
            torch.ones(3, 4, dtype=torch.long),
            "same batch size",
        ),
        (
            torch.zeros(2, 4, 8),
            torch.ones(2, 5, dtype=torch.long),
            "same sequence length",
        ),
    ],
)
def test_encode_prompt_validates_precomputed_embeds_and_mask_shapes(
    prompt_embeds,
    prompt_embeds_mask,
    match,
):
    pipeline = object.__new__(JoyImageEditPipeline)

    with pytest.raises(ValueError, match=match):
        JoyImageEditPipeline.encode_prompt(
            pipeline,
            "prompt",
            Image.new("RGB", (16, 16)),
            num_images_per_prompt=1,
            prompt_embeds=prompt_embeds,
            prompt_embeds_mask=prompt_embeds_mask,
        )


def test_encode_prompt_validation_uses_negative_parameter_names():
    pipeline = object.__new__(JoyImageEditPipeline)

    with pytest.raises(ValueError, match="negative_prompt_embeds_mask must be a 2D tensor"):
        JoyImageEditPipeline.encode_prompt(
            pipeline,
            "",
            Image.new("RGB", (16, 16)),
            num_images_per_prompt=1,
            prompt_embeds=torch.zeros(1, 4, 8),
            prompt_embeds_mask=torch.ones(1, 4, 1, dtype=torch.long),
            embeds_name="negative_prompt_embeds",
            mask_name="negative_prompt_embeds_mask",
        )


def test_last_layer_capture_passes_mm_token_type_ids():
    class FakeDecoderLayer(torch.nn.Module):
        def forward(self, hidden_states):
            return (hidden_states + 1.0,)

    class FakeLanguageModel(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = torch.nn.ModuleList([FakeDecoderLayer()])

    class FakeTextEncoder(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.model = SimpleNamespace(language_model=FakeLanguageModel())
            self.calls = []

        def forward(self, **kwargs):
            self.calls.append(kwargs)
            hidden_states = torch.zeros(1, 2, 4)
            self.model.language_model.layers[-1](hidden_states)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.text_encoder = FakeTextEncoder()
    mm_token_type_ids = torch.tensor([[0, 1]], dtype=torch.long)
    model_inputs = SimpleNamespace(
        input_ids=torch.tensor([[1, 2]], dtype=torch.long),
        attention_mask=torch.tensor([[1, 1]], dtype=torch.long),
        pixel_values=torch.zeros(1, 3, 16, 16),
        image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long),
        mm_token_type_ids=mm_token_type_ids,
        extra_processor_field=torch.tensor([7], dtype=torch.long),
    )

    hidden_states = JoyImageEditPipeline._get_last_layer_pre_norm_hidden(
        pipeline,
        model_inputs,
    )

    assert torch.equal(hidden_states, torch.ones(1, 2, 4))
    assert pipeline.text_encoder.calls[0]["mm_token_type_ids"] is mm_token_type_ids
    assert torch.equal(pipeline.text_encoder.calls[0]["extra_processor_field"], torch.tensor([7], dtype=torch.long))
    assert pipeline.text_encoder.calls[0]["output_hidden_states"] is False


def test_cast_floating_model_inputs_only_casts_pixel_tensors():
    class FakeModelInputs(dict):
        def __getattr__(self, name):
            return self[name]

    model_inputs = FakeModelInputs(
        input_ids=torch.tensor([[1, 2]], dtype=torch.long),
        attention_mask=torch.tensor([[1, 1]], dtype=torch.long),
        pixel_values=torch.zeros(1, 3, dtype=torch.float32),
        image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long),
    )

    _cast_floating_model_inputs(model_inputs, torch.bfloat16)

    assert model_inputs.input_ids.dtype == torch.long
    assert model_inputs.attention_mask.dtype == torch.long
    assert model_inputs.pixel_values.dtype == torch.bfloat16
    assert model_inputs.image_grid_thw.dtype == torch.long


def test_qwen_prompt_embeds_casts_processor_pixel_values_to_encoder_dtype():
    class FakeModelInputs(dict):
        def __init__(self):
            super().__init__(
                input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long),
                attention_mask=torch.tensor([[1, 1, 1]], dtype=torch.long),
                pixel_values=torch.zeros(1, 3, dtype=torch.float32),
                image_grid_thw=torch.tensor([[1, 1, 1]], dtype=torch.long),
            )

        def __getattr__(self, name):
            return self[name]

        def __setattr__(self, name, value):
            self[name] = value

        def to(self, device):
            for key, value in list(self.items()):
                if isinstance(value, torch.Tensor):
                    self[key] = value.to(device)
            return self

    class FakeProcessor:
        def __init__(self):
            self.model_inputs = FakeModelInputs()

        def __call__(self, **kwargs):
            self.kwargs = kwargs
            return self.model_inputs

    class FakeTextEncoder:
        dtype = torch.bfloat16

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.processor = FakeProcessor()
    pipeline.text_encoder = FakeTextEncoder()
    pipeline.prompt_template_encode = JOY_PROMPT_TEMPLATE
    pipeline.prompt_template_encode_start_idx = 0
    pipeline.tokenizer_max_length = 8
    captured = {}

    def fake_get_hidden(model_inputs):
        captured["pixel_dtype"] = model_inputs.pixel_values.dtype
        captured["attention_mask_dtype"] = model_inputs.attention_mask.dtype
        return torch.ones(1, 3, 4, dtype=torch.float32)

    pipeline._get_last_layer_pre_norm_hidden = fake_get_hidden

    prompt_embeds, prompt_mask = JoyImageEditPipeline._get_qwen_prompt_embeds(
        pipeline,
        "make it brighter",
        Image.new("RGB", (16, 16)),
    )

    assert captured == {
        "pixel_dtype": torch.bfloat16,
        "attention_mask_dtype": torch.long,
    }
    text = pipeline.processor.kwargs["text"][0]
    assert text.startswith("<|im_start|>system\n \\nDescribe")
    assert "Describe the image by detailing the color, shape, size, texture" in text
    assert (f"<|im_start|>user\n{JOY_VISION_TOKEN}make it brighter<|im_end|>\n") in text
    assert text.endswith("<|im_start|>assistant\n")
    assert prompt_embeds.dtype == torch.bfloat16
    assert prompt_mask.dtype == torch.long


def test_prepare_latents_casts_vae_encoded_latents_to_requested_dtype():
    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.latent_channels = 4
    pipeline.vae_scale_factor = 8

    def fake_encode(image, generator):
        assert image.dtype == torch.bfloat16
        return torch.ones(1, 4, 1, 8, 8, dtype=torch.float32)

    pipeline._encode_vae_image = fake_encode

    latents, image_latents = JoyImageEditPipeline._prepare_latents(
        pipeline,
        image=torch.zeros(1, 3, 1, 64, 64, dtype=torch.float32),
        batch_size=1,
        num_channels_latents=4,
        height=64,
        width=64,
        dtype=torch.bfloat16,
        device=torch.device("cpu"),
        generator=None,
    )

    assert latents.dtype == torch.bfloat16
    assert image_latents.dtype == torch.bfloat16


def test_prepare_latents_prepares_offloaded_vae_before_encode():
    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.vae_scale_factor = 8
    calls = []
    pipeline._prepare_vae_for_decode = lambda: calls.append("prepare")

    def encode(image, generator):
        calls.append("encode")
        return torch.zeros(1, 4, 1, 2, 2, device=image.device, dtype=image.dtype)

    pipeline._encode_vae_image = encode
    pipeline._prepare_latents(
        image=torch.zeros(1, 3, 1, 16, 16),
        batch_size=1,
        num_channels_latents=4,
        height=16,
        width=16,
        dtype=torch.float32,
        device=torch.device("cpu"),
        generator=None,
    )

    assert calls == ["prepare", "encode"]


def test_encode_vae_image_does_not_forward_noise_generator_to_posterior_sample():
    class FakeLatentDist:
        def __init__(self):
            self.received_generator = object()

        def sample(self, generator=None):
            self.received_generator = generator
            return torch.zeros(1, 4, 1, 8, 8)

    class FakeVAE:
        def __init__(self):
            self.config = SimpleNamespace(
                latents_mean=[0.0, 0.0, 0.0, 0.0],
                latents_std=[1.0, 1.0, 1.0, 1.0],
            )
            self.latent_dist = FakeLatentDist()

        def encode(self, image):
            return SimpleNamespace(latent_dist=self.latent_dist)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.vae = FakeVAE()
    pipeline.latent_channels = 4

    generator = torch.Generator(device="cpu").manual_seed(123)
    JoyImageEditPipeline._encode_vae_image(
        pipeline,
        torch.zeros(1, 3, 1, 64, 64),
        generator,
    )

    assert pipeline.vae.latent_dist.received_generator is None


@pytest.mark.parametrize(
    "dtype",
    [
        torch.float32,
        torch.bfloat16,
    ],
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_joy_attention_native_sdpa_matches_torch_sdpa_without_mask(dtype):
    torch.manual_seed(123)
    device = torch.device("cuda")
    attention = _make_joy_attention(dtype=dtype).to(device=device)
    hidden_states = torch.randn(2, 3, 32, device=device, dtype=dtype)
    encoder_hidden_states = torch.randn(2, 5, 32, device=device, dtype=dtype)

    actual = attention(hidden_states, encoder_hidden_states)
    expected = _reference_joy_attention_output(attention, hidden_states, encoder_hidden_states)

    assert isinstance(attention.attn, Attention)
    assert attention.attn.role == "joy_image.joint"
    assert attention.attn.role_category == "self"
    assert attention.attn.qkv_layout == "BSND"
    _assert_attention_outputs_close(actual, expected, dtype=dtype)


@pytest.mark.parametrize(
    "dtype",
    [
        torch.float32,
        torch.bfloat16,
    ],
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_joy_attention_native_sdpa_matches_torch_sdpa_with_padding_mask(dtype):
    torch.manual_seed(456)
    device = torch.device("cuda")
    attention = _make_joy_attention(dtype=dtype).to(device=device)
    hidden_states = torch.randn(2, 3, 32, device=device, dtype=dtype)
    encoder_hidden_states = torch.randn(2, 5, 32, device=device, dtype=dtype)
    attention_mask = torch.tensor(
        [
            [1, 1, 1, 1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1, 1, 0, 0],
        ],
        device=device,
        dtype=torch.bool,
    )

    actual = attention(hidden_states, encoder_hidden_states, attention_mask=attention_mask)
    expected = _reference_joy_attention_output(
        attention,
        hidden_states,
        encoder_hidden_states,
        attention_mask=attention_mask,
    )

    _assert_attention_outputs_close(actual, expected, dtype=dtype)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_transformer_shape_and_masked_forward():
    device = torch.device("cuda")
    transformer = JoyImageEditTransformer3DModel(
        in_channels=4,
        out_channels=4,
        hidden_size=32,
        text_dim=16,
        num_layers=1,
        num_attention_heads=4,
        patch_size=(1, 2, 2),
    ).to(device=device, dtype=torch.bfloat16)
    hidden_states = torch.randn(2, 2, 4, 1, 4, 4, device=device, dtype=torch.bfloat16)
    encoder_hidden_states = torch.randn(2, 5, 16, device=device, dtype=torch.bfloat16)
    encoder_hidden_states_mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 1, 0, 0]], device=device)

    output = transformer(
        hidden_states=hidden_states,
        timestep=torch.tensor([1.0, 2.0], device=device),
        encoder_hidden_states=encoder_hidden_states,
        encoder_hidden_states_mask=encoder_hidden_states_mask,
        return_dict=False,
    )[0]

    assert output.shape == hidden_states.shape
    assert isinstance(transformer.double_blocks[0].attn.attn, Attention)


def test_transformer_from_config_file_loads_checkpoint_config(tmp_path):
    _write_model_configs(tmp_path)

    transformer = JoyImageEditTransformer3DModel.from_config_file(tmp_path / "transformer" / "config.json")

    assert transformer.patch_size == (1, 2, 2)
    assert transformer.hidden_size == 32
    assert transformer.num_attention_heads == 4
    assert transformer.num_layers == 1
    assert transformer.text_dim == 16
    assert transformer.theta == 256
    assert len(transformer.double_blocks) == 1


def test_transformer_load_weights_maps_diffusers_prefix():
    transformer = JoyImageEditTransformer3DModel(
        in_channels=4,
        out_channels=4,
        hidden_size=32,
        text_dim=16,
        num_layers=1,
        num_attention_heads=4,
        patch_size=(1, 2, 2),
    )
    weight = torch.randn_like(transformer.img_in.weight)

    loaded = transformer.load_weights([("transformer.patch_embedding.weight", weight)])

    assert loaded == {"img_in.weight"}
    assert torch.equal(transformer.img_in.weight, weight)

    text_weight = torch.randn_like(transformer.condition_embedder.text_embedder.linear_1.weight)
    loaded = transformer.load_weights([("transformer.condition_embedder.text_embedder.linear_1.weight", text_weight)])

    assert loaded == {"condition_embedder.text_embedder.linear_1.weight"}
    assert torch.equal(transformer.condition_embedder.text_embedder.linear_1.weight, text_weight)


def test_transformer_load_weights_warns_for_unknown_weight(caplog):
    transformer = JoyImageEditTransformer3DModel(
        in_channels=4,
        out_channels=4,
        hidden_size=32,
        text_dim=16,
        num_layers=1,
        num_attention_heads=4,
        patch_size=(1, 2, 2),
    )

    caplog.set_level(logging.WARNING)
    loaded = transformer.load_weights([("transformer.unmapped.weight", torch.empty(1))])

    assert loaded == set()
    assert "Skipping JoyAI-Image-Edit transformer weight transformer.unmapped.weight" in caplog.text
    assert "checkpoint mismatch" in caplog.text


def test_latent_normalization_round_trip():
    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.latent_channels = 4
    pipeline.vae = SimpleNamespace(
        config=SimpleNamespace(
            latents_mean=[0.1, -0.2, 0.3, -0.4],
            latents_std=[0.5, 0.75, 1.25, 2.0],
        )
    )
    latents = torch.randn(2, 4, 1, 2, 2)

    latents_mean, latents_std = JoyImageEditPipeline._latent_stats(
        pipeline,
        latents.device,
        latents.dtype,
    )
    normalized = (latents - latents_mean) / latents_std

    assert torch.allclose(normalized * latents_std + latents_mean, latents)


def test_decode_latents_unnormalizes_and_selects_target_slot():
    class FakeVAE:
        config = SimpleNamespace(
            latents_mean=[0.1, -0.2, 0.3, -0.4],
            latents_std=[0.5, 0.75, 1.25, 2.0],
        )

        def decode(self, latents, return_dict=False):
            return (latents,)

    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.latent_channels = 4
    pipeline.vae = FakeVAE()
    latents = torch.zeros(1, 2, 4, 1, 2, 2)
    latents[:, -1] = 2.0

    decoded = JoyImageEditPipeline._decode_latents(pipeline, latents)

    latents_mean, latents_std = JoyImageEditPipeline._latent_stats(
        pipeline,
        latents.device,
        latents.dtype,
    )
    assert decoded.shape == (1, 4, 2, 2)
    assert torch.equal(decoded, (latents[:, -1] * latents_std + latents_mean)[:, :, 0])


def test_prepare_latents_stacks_reference_first_target_last():
    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.vae_scale_factor = 8
    pipeline.latent_channels = 4
    image_latents = torch.ones(1, 4, 1, 2, 2)
    noise_latents = torch.zeros(1, 1, 4, 1, 2, 2)
    pipeline._encode_vae_image = lambda image, generator: image_latents

    latents, reference_latents = JoyImageEditPipeline._prepare_latents(
        pipeline,
        image=torch.zeros(1, 3, 1, 16, 16),
        batch_size=1,
        num_channels_latents=4,
        height=16,
        width=16,
        dtype=torch.float32,
        device=torch.device("cpu"),
        generator=None,
        latents=noise_latents,
    )

    assert latents.shape == (1, 2, 4, 1, 2, 2)
    assert torch.equal(latents[:, :1], reference_latents)
    assert torch.equal(latents[:, -1:], noise_latents)


def test_prepare_latents_validates_generator_and_latent_shapes():
    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.vae_scale_factor = 8
    pipeline.latent_channels = 4
    image_latents = torch.ones(1, 4, 1, 2, 2)
    pipeline._encode_vae_image = lambda image, generator: image_latents

    with pytest.raises(ValueError, match="Generator list length"):
        JoyImageEditPipeline._prepare_latents(
            pipeline,
            image=torch.zeros(1, 3, 1, 16, 16),
            batch_size=2,
            num_channels_latents=4,
            height=16,
            width=16,
            dtype=torch.float32,
            device=torch.device("cpu"),
            generator=[torch.Generator()],
        )

    with pytest.raises(ValueError, match="noise latents must have shape"):
        JoyImageEditPipeline._prepare_latents(
            pipeline,
            image=torch.zeros(1, 3, 1, 16, 16),
            batch_size=1,
            num_channels_latents=4,
            height=16,
            width=16,
            dtype=torch.float32,
            device=torch.device("cpu"),
            generator=None,
            latents=torch.zeros(1, 1, 4, 1, 1, 2),
        )


def test_prepare_latents_validates_image_latent_shape():
    pipeline = object.__new__(JoyImageEditPipeline)
    pipeline.vae_scale_factor = 8
    pipeline.latent_channels = 4
    pipeline._encode_vae_image = lambda image, generator: torch.ones(1, 4, 2, 2)

    with pytest.raises(ValueError, match="image latents must have shape"):
        JoyImageEditPipeline._prepare_latents(
            pipeline,
            image=torch.zeros(1, 3, 1, 16, 16),
            batch_size=1,
            num_channels_latents=4,
            height=16,
            width=16,
            dtype=torch.float32,
            device=torch.device("cpu"),
            generator=None,
        )


def test_diffuse_restores_reference_slots_each_step():
    class FakeScheduler:
        def set_begin_index(self, index):
            self.begin_index = index

        def step(self, noise_pred, timestep, latents, return_dict=False):
            return (latents + 1.0,)

    class FakeTransformer(torch.nn.Module):
        def forward(self, hidden_states, **kwargs):
            return (torch.zeros_like(hidden_states),)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.scheduler = FakeScheduler()
    pipeline.transformer = FakeTransformer()
    pipeline._interrupt = False
    pipeline._current_timestep = None

    latents = torch.zeros(1, 2, 4, 1, 2, 2)
    image_latents = torch.full((1, 1, 4, 1, 2, 2), 7.0)
    result = JoyImageEditPipeline.diffuse(
        pipeline,
        latents=latents,
        image_latents=image_latents,
        prompt_embeds=torch.zeros(1, 3, 8),
        prompt_embeds_mask=torch.ones(1, 3, dtype=torch.long),
        negative_prompt_embeds=None,
        negative_prompt_embeds_mask=None,
        timesteps=torch.tensor([2.0, 1.0]),
        do_true_cfg=False,
        true_cfg_scale=1.0,
    )

    assert torch.equal(result[:, :1], image_latents)
    assert torch.equal(result[:, -1:], torch.full((1, 1, 4, 1, 2, 2), 2.0))


@pytest.mark.parametrize(
    "params_kwargs, expected_scale",
    [
        ({}, 4.0),
        ({"guidance_scale": 2.0}, 2.0),
        ({"guidance_scale": 1.0}, 1.0),
        ({"guidance_scale": 0.0}, 0.0),
        ({"true_cfg_scale": 8.0}, 8.0),
    ],
)
def test_forward_uses_guidance_scale_for_cfg(params_kwargs, expected_scale):
    class FakeScheduler:
        def set_timesteps(self, num_inference_steps, device):
            self.num_inference_steps = num_inference_steps
            self.timesteps = torch.tensor([1.0], device=device)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.latent_channels = 4
    pipeline.scheduler = FakeScheduler()
    encode_calls = []
    diffuse_calls = []

    def encode_prompt(prompt, image, **kwargs):
        encode_calls.append(prompt)
        return torch.zeros(1, 3, 4), torch.ones(1, 3, dtype=torch.long)

    def fake_prepare_latents(**kwargs):
        latents = torch.zeros(1, 2, 4, 1, 2, 2)
        image_latents = torch.ones(1, 1, 4, 1, 2, 2)
        return latents, image_latents

    def diffuse(**kwargs):
        diffuse_calls.append(kwargs)
        return kwargs["latents"]

    pipeline.encode_prompt = encode_prompt
    pipeline._prepare_latents = fake_prepare_latents
    pipeline.diffuse = diffuse
    pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 2, 2)

    request = _make_request(
        prompt={
            "prompt": "make it brighter",
            "additional_information": {
                "image_tensor": torch.zeros(1, 3, 1, 16, 16),
                "prompt_image": Image.new("RGB", (16, 16)),
                "height": 16,
                "width": 16,
            },
        },
        params=_make_params(**params_kwargs),
    )

    output = JoyImageEditPipeline.forward(pipeline, request)

    assert output.output.shape == (1, 3, 2, 2)
    expected_cfg = expected_scale > 1.0
    assert encode_calls == (["make it brighter", ""] if expected_cfg else ["make it brighter"])
    assert diffuse_calls[0]["do_true_cfg"] is expected_cfg
    assert diffuse_calls[0]["true_cfg_scale"] == expected_scale


@pytest.mark.parametrize(
    "schedule_kwargs, expected_timesteps",
    [
        ({"timesteps": torch.tensor([9, 4])}, torch.tensor([9.0, 4.0])),
        ({"sigmas": [1.0, 0.5]}, torch.tensor([1000.0, 500.0])),
    ],
)
def test_forward_uses_request_sampling_parameters(schedule_kwargs, expected_timesteps):
    class FakeScheduler:
        def set_timesteps(self, num_inference_steps=None, device=None, timesteps=None, sigmas=None):
            if timesteps is not None:
                self.timesteps = torch.as_tensor(timesteps, device=device, dtype=torch.float32)
            elif sigmas is not None:
                self.timesteps = torch.as_tensor(sigmas, device=device) * 1000
            else:
                self.timesteps = torch.arange(num_inference_steps, device=device, dtype=torch.float32)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.latent_channels = 4
    pipeline.tokenizer_max_length = JOY_MAX_IMAGE_SEQ_LEN
    pipeline.scheduler = FakeScheduler()
    encode_calls = []
    prepare_calls = []
    diffuse_calls = []

    def encode_prompt(prompt, image, **kwargs):
        encode_calls.append(kwargs)
        return torch.zeros(1, 3, 4), torch.ones(1, 3, dtype=torch.long)

    def prepare_latents(**kwargs):
        prepare_calls.append(kwargs)
        reference = torch.ones(1, 1, 4, 1, 2, 2)
        return torch.cat([reference, kwargs["latents"]], dim=1), reference

    def diffuse(**kwargs):
        diffuse_calls.append(kwargs)
        return kwargs["latents"]

    pipeline.encode_prompt = encode_prompt
    pipeline._prepare_latents = prepare_latents
    pipeline.diffuse = diffuse
    pipeline._decode_latents = lambda _latents: pytest.fail("latent output must not be decoded")

    request_latents = torch.full((1, 1, 4, 1, 2, 2), 3.0)
    params = _make_params(
        guidance_scale=1.0,
        latents=request_latents,
        output_type="latent",
        max_sequence_length=123,
        **schedule_kwargs,
    )
    request = _make_request(
        prompt={
            "prompt": "make it brighter",
            "additional_information": {
                "image_tensor": torch.zeros(1, 3, 1, 16, 16),
                "prompt_image": Image.new("RGB", (16, 16)),
                "height": 16,
                "width": 16,
            },
        },
        params=params,
    )

    output = JoyImageEditPipeline.forward(pipeline, request)

    assert encode_calls[0]["max_sequence_length"] == 123
    assert prepare_calls[0]["latents"] is request_latents
    torch.testing.assert_close(diffuse_calls[0]["timesteps"], expected_timesteps)
    assert diffuse_calls[0]["cfg_normalize"] is True
    torch.testing.assert_close(output.output, request_latents[:, 0])


def test_retrieve_timesteps_rejects_conflicting_custom_schedules():
    scheduler = SimpleNamespace(set_timesteps=lambda **_kwargs: None)

    with pytest.raises(ValueError, match="Only one"):
        retrieve_timesteps(scheduler, timesteps=[9], sigmas=[1.0])


def test_forward_rejects_unsupported_output_type():
    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    request = _make_request(
        prompt={
            "prompt": "make it brighter",
            "additional_information": {
                "image_tensor": torch.zeros(1, 3, 1, 16, 16),
                "prompt_image": Image.new("RGB", (16, 16)),
                "height": 16,
                "width": 16,
            },
        },
        params=_make_params(output_type="np"),
    )

    with pytest.raises(ValueError, match="only `pil` and `latent`"):
        JoyImageEditPipeline.forward(pipeline, request)


def test_forward_skips_decode_for_dummy_warmup_request(monkeypatch):
    class FakeScheduler:
        def set_timesteps(self, num_inference_steps, device):
            self.timesteps = torch.tensor([1.0], device=device)

    pipeline = object.__new__(JoyImageEditPipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.device = torch.device("cpu")
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=True,
        enable_layerwise_offload=False,
        pin_cpu_memory=False,
        parallel_config=SimpleNamespace(use_hsdp=False),
    )
    pipeline.transformer = torch.nn.Linear(1, 1)
    pipeline.latent_channels = 4
    pipeline.scheduler = FakeScheduler()
    pipeline.encode_prompt = lambda *args, **kwargs: (
        torch.zeros(1, 3, 4),
        torch.ones(1, 3, dtype=torch.long),
    )
    pipeline._prepare_latents = lambda **kwargs: (
        torch.zeros(1, 2, 4, 1, 2, 2),
        torch.ones(1, 1, 4, 1, 2, 2),
    )
    pipeline.diffuse = lambda **kwargs: kwargs["latents"] + 1.0
    pipeline.check_cfg_parallel_validity = lambda scale, has_neg_prompt: True

    def decode_should_not_run(latents):
        raise AssertionError("dummy warmup should not decode latents")

    pipeline._decode_latents = decode_should_not_run
    move_calls = []

    def fake_move_params(module, target_device, **kwargs):
        move_calls.append((module, target_device, kwargs))

    monkeypatch.setattr(
        joy_pipeline_module.SequentialOffloadHook,
        "_move_params",
        fake_move_params,
    )
    monkeypatch.setattr(joy_pipeline_module.current_omni_platform, "empty_cache", lambda: None)

    request = _make_request(
        prompt={
            "prompt": "dummy run",
            "additional_information": {
                "image_tensor": torch.zeros(1, 3, 1, 16, 16),
                "prompt_image": Image.new("RGB", (16, 16)),
                "height": 16,
                "width": 16,
            },
        },
        params=_make_params(guidance_scale=1.0),
    )
    request.request_id = DUMMY_DIFFUSION_REQUEST_ID
    request.is_dummy_run = lambda: True

    output = JoyImageEditPipeline.forward(pipeline, request)

    assert output.output.shape == (1, 4, 1, 2, 2)
    assert move_calls == [
        (
            pipeline.transformer,
            torch.device("cpu"),
            {"non_blocking": False, "pin_memory": False},
        )
    ]