# 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}, ) ]