diff --git a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py index 5d608d7c49fb..99733ccc3c08 100644 --- a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py +++ b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py @@ -344,9 +344,12 @@ def __call__( ) # Start from a fully random canvas and denoise it; the scheduler resets its committed state at step 0. + # `torch.randint` requires the generator and the output device to match, so (as with `randn_tensor`) a + # CPU generator samples on CPU and the result is moved to `device` afterwards. + rand_device = generator.device if generator is not None else device canvas = torch.randint( - 0, text_config.vocab_size, (batch_size, canvas_length), device=device, generator=generator - ) + 0, text_config.vocab_size, (batch_size, canvas_length), device=rand_device, generator=generator + ).to(device) self_conditioning_logits = None finished_denoising = torch.zeros(batch_size, dtype=torch.bool, device=device) argmax_canvas = canvas diff --git a/src/diffusers/pipelines/llada2/pipeline_llada2.py b/src/diffusers/pipelines/llada2/pipeline_llada2.py index 06b4875f18a9..e6128d8a6ec3 100644 --- a/src/diffusers/pipelines/llada2/pipeline_llada2.py +++ b/src/diffusers/pipelines/llada2/pipeline_llada2.py @@ -71,6 +71,8 @@ class LLaDA2Pipeline(DiffusionPipeline): scheduler: BlockRefinementScheduler tokenizer: Any + _optional_components = ["tokenizer"] + _callback_tensor_inputs = [ "block_x", "transfer_index", diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 69bce1a1c533..e13d6dbea4a6 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -936,7 +936,11 @@ def _fetch_class_library_tuple(module): pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None path = not_compiled_module.__module__.split(".") - is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) + # A same-named folder in another library (e.g. `transformers.models.diffusion_gemma` vs + # `diffusers.pipelines.diffusion_gemma`) must not count as a pipeline module. + is_pipeline_module = ( + path[0] == diffusers_module.__name__ and pipeline_dir in path and hasattr(pipelines, pipeline_dir) + ) # if library is not in LOADABLE_CLASSES, then it is a custom module. # Or if it's a pipeline module, then the module is inside the pipeline diff --git a/src/diffusers/schedulers/scheduling_block_refinement.py b/src/diffusers/schedulers/scheduling_block_refinement.py index 6ff7963748b0..739c9e771098 100644 --- a/src/diffusers/schedulers/scheduling_block_refinement.py +++ b/src/diffusers/schedulers/scheduling_block_refinement.py @@ -173,7 +173,10 @@ def _sample_from_logits( filtered = BlockRefinementScheduler._top_p_filtering(filtered, top_p=top_p) probs = torch.softmax(filtered.float(), dim=-1) - token = torch.multinomial(probs, num_samples=1, generator=generator) + # `torch.multinomial` requires the generator and the sampled tensor's device to match, so (as with + # `randn_tensor`) a CPU generator samples on CPU and the result is moved back to `probs`'s device. + rand_device = generator.device if generator is not None else probs.device + token = torch.multinomial(probs.to(rand_device), num_samples=1, generator=generator).to(probs.device) token_prob = torch.gather(probs, -1, token) return token.view(*logits.shape[:-1]), token_prob.view(*logits.shape[:-1]) @@ -285,9 +288,10 @@ def step( prev_sample = torch.where(transfer_index | editing_transfer_index, sampled_tokens, sample) self._committed = committed | transfer_index + rand_device = generator.device if generator is not None else sample.device random_tokens = torch.randint( - low=0, high=model_output.shape[-1], size=sample.shape, device=sample.device, generator=generator - ) + low=0, high=model_output.shape[-1], size=sample.shape, device=rand_device, generator=generator + ).to(sample.device) prev_sample = torch.where(self._committed, prev_sample, random_tokens) if not return_dict: @@ -498,14 +502,15 @@ def add_noise( masked_rev = torch.zeros_like(original_samples, dtype=torch.bool) valid = attention_mask.to(dtype=torch.bool) + rand_device = generator.device if generator is not None else device for block_start in range(prompt_length, seq_len, block_length): block_end = min(seq_len, block_start + block_length) seg_len = block_end - block_start if seg_len <= 0: continue - p_mask = torch.rand((batch_size, 1), device=device, generator=generator) - seg = torch.rand((batch_size, seg_len), device=device, generator=generator) < p_mask + p_mask = torch.rand((batch_size, 1), device=rand_device, generator=generator).to(device) + seg = torch.rand((batch_size, seg_len), device=rand_device, generator=generator).to(device) < p_mask seg = seg & valid[:, block_start:block_end] seg_rev = (~seg) & valid[:, block_start:block_end] diff --git a/src/diffusers/schedulers/scheduling_discrete_ddim.py b/src/diffusers/schedulers/scheduling_discrete_ddim.py index fff98edced00..3fa598c93321 100644 --- a/src/diffusers/schedulers/scheduling_discrete_ddim.py +++ b/src/diffusers/schedulers/scheduling_discrete_ddim.py @@ -117,7 +117,12 @@ def _sample_from_logits( token = flat_logits.argmax(dim=-1, keepdim=True) else: scaled_probs = torch.softmax(flat_logits.float() / temperature, dim=-1) - token = torch.multinomial(scaled_probs, num_samples=1, generator=generator) + # `torch.multinomial` requires the generator and the sampled tensor's device to match, so (as with + # `randn_tensor`) a CPU generator samples on CPU and the result is moved back to `scaled_probs`'s device. + rand_device = generator.device if generator is not None else scaled_probs.device + token = torch.multinomial(scaled_probs.to(rand_device), num_samples=1, generator=generator).to( + scaled_probs.device + ) token_prob = torch.gather(probs, -1, token) return token.view(*logits.shape[:-1]), token_prob.view(*logits.shape[:-1]) @@ -200,11 +205,16 @@ def step( route_probs = torch.stack([clean_mass, stay_mass, noise_mass], dim=-1) route_probs = route_probs / route_probs.sum(dim=-1, keepdim=True) - routes = torch.multinomial(route_probs.view(-1, 3), num_samples=1, generator=generator).view_as(sample) + rand_device = generator.device if generator is not None else route_probs.device + routes = ( + torch.multinomial(route_probs.view(-1, 3).to(rand_device), num_samples=1, generator=generator) + .to(route_probs.device) + .view_as(sample) + ) random_tokens = torch.randint( - low=0, high=vocab_size, size=sample.shape, device=sample.device, generator=generator - ) + low=0, high=vocab_size, size=sample.shape, device=rand_device, generator=generator + ).to(sample.device) prev_sample = torch.where(routes == 0, sampled_tokens, sample) prev_sample = torch.where(routes == 2, random_tokens, prev_sample) @@ -226,7 +236,8 @@ def _select_positions( k_eff = min(max(1, int(self.config.corrector_k)), seq_len) if selection == "random": - scores = torch.rand(batch_size, seq_len, device=sample.device, generator=generator) + rand_device = generator.device if generator is not None else sample.device + scores = torch.rand(batch_size, seq_len, device=rand_device, generator=generator).to(sample.device) return torch.topk(scores, k=k_eff, dim=-1).indices if selection == "lowest_maxprob": @@ -241,7 +252,8 @@ def _select_positions( raise ValueError(f"Unknown `corrector_selection`: {selection!r}.") keys = confidence / float(self.config.corrector_selection_tau) - u = torch.rand(keys.shape, device=keys.device, generator=generator).clamp_(1e-12, 1.0 - 1e-12) + rand_device = generator.device if generator is not None else keys.device + u = torch.rand(keys.shape, device=rand_device, generator=generator).to(keys.device).clamp_(1e-12, 1.0 - 1e-12) keys = keys + (-torch.log(-torch.log(u))) return torch.topk(keys, k=k_eff, dim=-1).indices @@ -295,9 +307,12 @@ def step_correct( positions = self._select_positions(sample, cond_log_probs, generator) rows = torch.arange(sample.shape[0], device=sample.device).unsqueeze(-1).expand_as(positions) chosen_probs = cond_log_probs[rows, positions].exp() - resampled = torch.multinomial( - chosen_probs.reshape(-1, vocab_size), num_samples=1, generator=generator - ).view_as(positions) + rand_device = generator.device if generator is not None else chosen_probs.device + resampled = ( + torch.multinomial(chosen_probs.reshape(-1, vocab_size).to(rand_device), num_samples=1, generator=generator) + .to(chosen_probs.device) + .view_as(positions) + ) prev_sample = sample.clone() prev_sample[rows, positions] = resampled diff --git a/src/diffusers/schedulers/scheduling_entropy_bound.py b/src/diffusers/schedulers/scheduling_entropy_bound.py index 5382190ec6bf..f59c81710b12 100644 --- a/src/diffusers/schedulers/scheduling_entropy_bound.py +++ b/src/diffusers/schedulers/scheduling_entropy_bound.py @@ -110,7 +110,12 @@ def _sample_from_logits( token = flat_logits.argmax(dim=-1, keepdim=True) else: scaled_probs = torch.softmax(flat_logits.float() / temperature, dim=-1) - token = torch.multinomial(scaled_probs, num_samples=1, generator=generator) + # `torch.multinomial` requires the generator and the sampled tensor's device to match, so (as with + # `randn_tensor`) a CPU generator samples on CPU and the result is moved back to `scaled_probs`'s device. + rand_device = generator.device if generator is not None else scaled_probs.device + token = torch.multinomial(scaled_probs.to(rand_device), num_samples=1, generator=generator).to( + scaled_probs.device + ) token_prob = torch.gather(probs, -1, token) return token.view(*logits.shape[:-1]), token_prob.view(*logits.shape[:-1]) @@ -166,9 +171,10 @@ def step( input=torch.zeros_like(sorted_accepted), dim=-1, index=sorted_indices, src=sorted_accepted ) + rand_device = generator.device if generator is not None else sample.device random_tokens = torch.randint( - low=0, high=model_output.shape[-1], size=sample.shape, device=sample.device, generator=generator - ) + low=0, high=model_output.shape[-1], size=sample.shape, device=rand_device, generator=generator + ).to(sample.device) prev_sample = torch.where(accepted_index, sampled_tokens, random_tokens) if not return_dict: diff --git a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py index b7ccb3e9d91d..63fb8fbdd6b3 100644 --- a/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py +++ b/tests/pipelines/diffusion_gemma/test_diffusion_gemma.py @@ -13,16 +13,21 @@ ) from diffusers.utils.import_utils import is_peft_available -from ...testing_utils import require_peft_backend, require_peft_version_greater +from ...testing_utils import ( + enable_full_determinism, + require_peft_backend, + require_peft_version_greater, + torch_device, +) +from ..pipeline_params import TEXT_TO_TEXT_BATCH_PARAMS, TEXT_TO_TEXT_PARAMS +from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin if is_peft_available(): from peft import LoraConfig -# `DiffusionGemmaPipeline` is a discrete *text* diffusion pipeline: it returns token sequences rather than images, -# so the image/video oriented `BasePipelineTesterConfig` + `PipelineTesterMixin` contract in `..testing_utils` -# does not apply here. These are plain pytest classes instead. +enable_full_determinism() # --- Lightweight stand-in for input-validation tests that never reach the model --- @@ -78,28 +83,67 @@ def test_prompt_and_messages_together_raises(self): _MODEL_ID = "trl-internal-testing/tiny-DiffusionGemmaForBlockDiffusion" -def _load_pipeline(): - try: - from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion - except ImportError as e: - pytest.skip(f"transformers without DiffusionGemma: {e}") - try: - model = DiffusionGemmaForBlockDiffusion.from_pretrained(_MODEL_ID, dtype=torch.float32).eval() - processor = AutoProcessor.from_pretrained(_MODEL_ID) - except Exception as e: # noqa: BLE001 - offline / hub errors should skip, not fail - pytest.skip(f"tiny DiffusionGemma checkpoint unavailable: {e}") - pipe = DiffusionGemmaPipeline(model=model, scheduler=BlockRefinementScheduler(), processor=processor) - pipe.set_progress_bar_config(disable=True) - return pipe, model.config.canvas_length +class DiffusionGemmaPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = DiffusionGemmaPipeline + required_input_params_in_call_signature = TEXT_TO_TEXT_PARAMS + batch_input_params = TEXT_TO_TEXT_BATCH_PARAMS + # DiffusionGemma has neither `num_images_per_prompt` (batching is over `prompt` only) nor `latents` (each + # canvas is freshly randomized inside `__call__`), and `output_type` is `"seq"`/`"text"`, not an image format. + optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"]) + # One canvas' worth of generated tokens for the tiny checkpoint's `canvas_length` (see `get_dummy_inputs`). + output_shape = (32,) + + def get_dummy_components(self): + try: + from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion + except ImportError as e: + pytest.skip(f"transformers without DiffusionGemma: {e}") + try: + model = DiffusionGemmaForBlockDiffusion.from_pretrained(_MODEL_ID, dtype=torch.float32).eval() + processor = AutoProcessor.from_pretrained(_MODEL_ID) + except Exception as e: # noqa: BLE001 - offline / hub errors should skip, not fail + pytest.skip(f"tiny DiffusionGemma checkpoint unavailable: {e}") + return {"model": model, "scheduler": BlockRefinementScheduler(), "processor": processor} + + def get_dummy_inputs(self): + return { + "prompt": "Name a color.", + "generator": self.get_generator(0), + "gen_length": self.output_shape[0], + "num_inference_steps": 4, + "temperature": 0.0, + "eos_early_stop": False, + "output_type": "seq", + } -class TestDiffusionGemmaPipeline: +class TestDiffusionGemmaPipeline(DiffusionGemmaPipelineTesterConfig, PipelineTesterMixin): adaptive_stopping_vocab_size = 8 prompt = "Name a color." @pytest.fixture(autouse=True) def pipeline(self): - self.pipe, self.canvas_length = _load_pipeline() + self.pipe = self.get_pipeline().to(torch_device) + self.canvas_length = self.pipe.model.config.canvas_length + + # DiffusionGemma samples its canvas with a single `torch.randint(..., generator=generator)` call, unlike the + # `randn_tensor`-backed image pipelines the base test assumes, so it can't take a per-batch-row generator list. + def test_inference_batch_consistent(self): + super().test_inference_batch_consistent(batch_generator=False) + + @pytest.mark.skip( + "Test not supported: passes a per-row generator list, which DiffusionGemma's single `torch.randint` " + "canvas init doesn't accept." + ) + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip( + "Test not supported: assumes an image/video pipeline (`output_type='latent'`, tensor key `'latents'`), " + "neither of which DiffusionGemma's `check_inputs` accepts." + ) + def test_callback_inputs(self): + pass def _run_adaptive_stopping(self, prompt): self.pipe.model.config.get_text_config(decoder=True).vocab_size = self.adaptive_stopping_vocab_size diff --git a/tests/pipelines/llada2/test_llada2.py b/tests/pipelines/llada2/test_llada2.py index 33b634a7e11a..588ad57a09df 100644 --- a/tests/pipelines/llada2/test_llada2.py +++ b/tests/pipelines/llada2/test_llada2.py @@ -1,8 +1,16 @@ import pytest import torch +from transformers import CLIPTokenizer, GPT2Config, GPT2LMHeadModel from diffusers import BlockRefinementScheduler, LLaDA2Pipeline +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..pipeline_params import TEXT_TO_TEXT_BATCH_PARAMS, TEXT_TO_TEXT_PARAMS +from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin + + +enable_full_determinism() + class _DummyModelOutput: def __init__(self, logits): @@ -40,7 +48,92 @@ def _make_pipeline(tokenizer=None): return LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer) -class TestLLaDA2Pipeline: +class LLaDA2PipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = LLaDA2Pipeline + required_input_params_in_call_signature = TEXT_TO_TEXT_PARAMS + batch_input_params = TEXT_TO_TEXT_BATCH_PARAMS + # LLaDA2 has neither `num_images_per_prompt` (batching is over `prompt` only) nor `latents` (the template is a + # fully-masked sequence built fresh inside `__call__`), and `output_type` is `"seq"`/`"text"`, not an image format. + optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"]) + # `gen_length` for the dummy inputs below (see `get_dummy_inputs`); `_DummyCausalLM` ignores token values (only + # reads `input_ids.shape`), so this is independent of the tokenizer's own vocab size. + output_shape = (16,) + mask_token_id = 31 + + def get_dummy_components(self): + # `LLaDA2Pipeline.model` accepts any object exposing `forward(input_ids, attention_mask, position_ids) -> + # logits`, so unlike `DiffusionGemma`'s VLM-backed pipeline there's no pretrained checkpoint to pull. The + # `save_pretrained`/`from_pretrained` round trip the mixin exercises does need a real `PreTrainedModel` + # though (the `_DummyCausalLM` stand-in the hand-written tests below use isn't one), so build a tiny + # `GPT2LMHeadModel` locally instead, matching its `vocab_size` to the tokenizer's like other pipeline tests do. + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + torch.manual_seed(0) + model = GPT2LMHeadModel(GPT2Config(n_embd=16, n_head=1, n_layer=1, vocab_size=len(tokenizer), n_ctx=99)) + return {"model": model, "scheduler": BlockRefinementScheduler(), "tokenizer": tokenizer} + + def get_dummy_inputs(self): + return {"prompt": "Name a color.", **self._common_inputs()} + + def _common_inputs(self): + """Generation knobs shared by every dummy call, regardless of how the prompt is supplied.""" + return { + "use_chat_template": False, + "generator": self.get_generator(0), + "gen_length": self.output_shape[0], + "block_length": self.output_shape[0], + "num_inference_steps": 4, + "temperature": 0.0, + "threshold": 2.0, # force top-k commits so every step transfers a deterministic number of tokens + "minimal_topk": 1, + "eos_early_stop": False, + "mask_token_id": self.mask_token_id, + "output_type": "seq", + } + + +class TestLLaDA2Pipeline(LLaDA2PipelineTesterConfig, PipelineTesterMixin): + # LLaDA2 samples its template with a single `torch.randint`/`torch.multinomial` call inside the scheduler, unlike + # the `randn_tensor`-backed image pipelines the base test assumes, so it can't take a per-batch-row generator list. + def test_inference_batch_consistent(self): + super().test_inference_batch_consistent(batch_generator=False) + + @pytest.mark.skip( + "Test not supported: passes a per-row generator list, which the scheduler's single-generator sampling " + "doesn't accept." + ) + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip( + "Test not supported: assumes an image/video pipeline (`output_type='latent'`, tensor key `'latents'`), " + "neither of which LLaDA2's `check_inputs` accepts." + ) + def test_callback_inputs(self): + pass + + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + # Adapted from the base test: dropping `tokenizer` means there's nothing left to encode a `prompt` string + # with, so the dummy input switches to pre-tokenized `input_ids` instead (see `_common_inputs`). + pipe = self.get_pipeline().to(torch_device) + pipe.tokenizer = None + + input_ids = torch.tensor([[5, 6, 7, 8]], dtype=torch.long) + output = pipe(input_ids=input_ids, **self._common_inputs())[0] + + pipe.save_pretrained(tmp_path, safe_serialization=False) + pipe_loaded = self.pipeline_class.from_pretrained(tmp_path) + pipe_loaded.to(torch_device) + pipe_loaded.set_progress_bar_config(disable=None) + assert pipe_loaded.tokenizer is None, "`tokenizer` did not stay set to None after loading." + + output_loaded = pipe_loaded(input_ids=input_ids, **self._common_inputs())[0] + assert_tensors_close( + output_loaded, + output, + atol=expected_max_difference, + msg="Output changed after dropping optional components.", + ) + def test_pipeline_runs(self): pipe = _make_pipeline().to("cpu") diff --git a/tests/pipelines/pipeline_params.py b/tests/pipelines/pipeline_params.py index 3db7c9fa1b0c..1ca87d59ef9f 100644 --- a/tests/pipelines/pipeline_params.py +++ b/tests/pipelines/pipeline_params.py @@ -101,6 +101,8 @@ UNCONDITIONAL_AUDIO_GENERATION_PARAMS = frozenset(["batch_size"]) +TEXT_TO_TEXT_PARAMS = frozenset(["prompt", "gen_length", "num_inference_steps"]) + # image params TEXT_TO_IMAGE_IMAGE_PARAMS = frozenset([]) @@ -130,5 +132,7 @@ VIDEO_TO_VIDEO_BATCH_PARAMS = frozenset(["prompt", "negative_prompt", "video"]) +TEXT_TO_TEXT_BATCH_PARAMS = frozenset(["prompt"]) + # callback params TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS = frozenset(["prompt_embeds"]) diff --git a/tests/pipelines/test_pipeline_utils.py b/tests/pipelines/test_pipeline_utils.py index c62e582597fc..0d140c13adf8 100644 --- a/tests/pipelines/test_pipeline_utils.py +++ b/tests/pipelines/test_pipeline_utils.py @@ -1100,3 +1100,29 @@ def test_push_to_hub_library_name(self): # Reset repo delete_repo(repo_id, token=TOKEN) + + +class TestFetchClassLibraryTuple: + def test_diffusers_model(self): + from diffusers import UNet2DConditionModel + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + assert _fetch_class_library_tuple(UNet2DConditionModel) == ("diffusers", "UNet2DConditionModel") + + def test_pipeline_module_class(self): + from diffusers.pipelines.deepfloyd_if import IFWatermarker + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + assert _fetch_class_library_tuple(IFWatermarker) == ("deepfloyd_if", "IFWatermarker") + + def test_other_library_class_shadowing_pipeline_dir(self): + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + # A transformers class whose model folder shares its name with a diffusers pipeline folder + # (e.g. `transformers.models.diffusion_gemma` vs `diffusers.pipelines.diffusion_gemma`) must + # resolve to its own library, not to the pipeline folder. + class FakeModel: + pass + + FakeModel.__module__ = "transformers.models.diffusion_gemma.modeling_diffusion_gemma" + assert _fetch_class_library_tuple(FakeModel) == ("transformers", "FakeModel")