From 80a11ac3d54171c1be969c17a871b32089f74ed5 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Sun, 30 Aug 2026 14:45:14 -0700 Subject: [PATCH 1/9] fix component library resolution when a transformers model folder shadows a pipeline dir --- .../pipelines/pipeline_loading_utils.py | 6 ++++- tests/pipelines/test_pipeline_utils.py | 26 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) 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/tests/pipelines/test_pipeline_utils.py b/tests/pipelines/test_pipeline_utils.py index 6bf79ae5dde3..521dc8802d73 100644 --- a/tests/pipelines/test_pipeline_utils.py +++ b/tests/pipelines/test_pipeline_utils.py @@ -1111,3 +1111,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") From f6d6f137b34f2cdbffbb4e5b95910189cf7f08cf Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:35:17 +0200 Subject: [PATCH 2/9] fix generator/device mismatch in discrete schedulers and DiffusionGemma's canvas init torch.randint/multinomial/rand need the generator and the sampled tensor on the same device. A CPU generator (the portable default recommended for reproducible pipeline calls) broke on any accelerator. --- .../pipeline_diffusion_gemma.py | 7 ++-- .../schedulers/scheduling_block_refinement.py | 12 ++++--- .../schedulers/scheduling_discrete_ddim.py | 35 +++++++++++++------ .../schedulers/scheduling_entropy_bound.py | 14 +++++--- src/diffusers/schedulers/scheduling_utils.py | 9 +++++ 5 files changed, 57 insertions(+), 20 deletions(-) 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/schedulers/scheduling_block_refinement.py b/src/diffusers/schedulers/scheduling_block_refinement.py index 6ff7963748b0..0bd462d210bc 100644 --- a/src/diffusers/schedulers/scheduling_block_refinement.py +++ b/src/diffusers/schedulers/scheduling_block_refinement.py @@ -20,7 +20,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin +from .scheduling_utils import SchedulerMixin, _generator_device @dataclass @@ -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; `_generator_device` + # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. + rand_device = _generator_device(probs, generator) + 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(sample, generator) 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: diff --git a/src/diffusers/schedulers/scheduling_discrete_ddim.py b/src/diffusers/schedulers/scheduling_discrete_ddim.py index fff98edced00..96f4d22cb21b 100644 --- a/src/diffusers/schedulers/scheduling_discrete_ddim.py +++ b/src/diffusers/schedulers/scheduling_discrete_ddim.py @@ -21,7 +21,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin +from .scheduling_utils import SchedulerMixin, _generator_device @dataclass @@ -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; `_generator_device` + # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. + rand_device = _generator_device(scaled_probs, generator) + 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(route_probs, generator) + 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(sample, generator) + 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(keys, generator) + 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(chosen_probs, generator) + 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..4c39263ba518 100644 --- a/src/diffusers/schedulers/scheduling_entropy_bound.py +++ b/src/diffusers/schedulers/scheduling_entropy_bound.py @@ -20,7 +20,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin +from .scheduling_utils import SchedulerMixin, _generator_device @dataclass @@ -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; `_generator_device` + # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. + rand_device = _generator_device(scaled_probs, generator) + 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(sample, generator) 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/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py index 8cdb21c2f011..1d191a2dbbfc 100644 --- a/src/diffusers/schedulers/scheduling_utils.py +++ b/src/diffusers/schedulers/scheduling_utils.py @@ -27,6 +27,15 @@ SCHEDULER_CONFIG_NAME = "scheduler_config.json" +def _generator_device(tensor: torch.Tensor, generator: torch.Generator | None) -> torch.device: + """Device to sample on for a discrete (`torch.randint`/`torch.multinomial`/`torch.rand`) draw: the generator's + own device when one is passed, else `tensor`'s device. Mirrors `randn_tensor`'s CPU-generator portability (a CPU + generator, the recommended default for reproducible pipeline calls, samples on CPU regardless of `tensor`'s device) + for the discrete samplers, which have no `randn_tensor` equivalent to call into. + """ + return generator.device if generator is not None else tensor.device + + # NOTE: We make this type an enum because it simplifies usage in docs and prevents # circular imports when used for `_compatibles` within the schedulers module. # When it's used as a type in pipelines, it really is a Union because the actual From 897c14323ed410871db8d58ec23ca1d1b5a51dd2 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:35:24 +0200 Subject: [PATCH 3/9] mark LLaDA2's tokenizer as an optional component it already defaults to None and works without one; just wasn't declared, so save/load silently dropped it. --- src/diffusers/pipelines/llada2/pipeline_llada2.py | 2 ++ 1 file changed, 2 insertions(+) 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", From 61df97b55e69214628f423d623fd17c6d02af18e Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:35:28 +0200 Subject: [PATCH 4/9] run DiffusionGemma pipeline tests through the shared PipelineTesterMixin follows the audioldm2 pattern for a non-image output; skips the handful of tests that assume guidance/latents/generator-lists. --- .../diffusion_gemma/test_diffusion_gemma.py | 82 ++++++++++++++----- tests/pipelines/pipeline_params.py | 4 + 2 files changed, 67 insertions(+), 19 deletions(-) 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/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"]) From 640e064d751afa759ac34e73877a917581fe51e6 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:35:34 +0200 Subject: [PATCH 5/9] run LLaDA2 pipeline tests through the shared PipelineTesterMixin same treatment as DiffusionGemma; keeps the existing regression tests as-is alongside it. --- tests/pipelines/llada2/test_llada2.py | 158 +++++++++++++++++++------- 1 file changed, 114 insertions(+), 44 deletions(-) diff --git a/tests/pipelines/llada2/test_llada2.py b/tests/pipelines/llada2/test_llada2.py index 6b00e133c7b1..c398e1dc7ece 100644 --- a/tests/pipelines/llada2/test_llada2.py +++ b/tests/pipelines/llada2/test_llada2.py @@ -1,9 +1,16 @@ -import unittest - +import pytest import torch +from transformers import CLIPTokenizer, GPT2Config, GPT2LMHeadModel from diffusers import BlockRefinementScheduler, LLaDA2Pipeline +from ...testing_utils import enable_full_determinism +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): @@ -41,7 +48,74 @@ def _make_pipeline(tokenizer=None): return LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer) -class LLaDA2PipelineTest(unittest.TestCase): +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.", + "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 + + @pytest.mark.skip( + "Test not supported: drops `tokenizer` and reruns the dummy `prompt` input, but dropping the tokenizer " + "means the caller must switch to pre-tokenized `input_ids` instead (see " + "`test_output_type_text_without_tokenizer` below, which covers this)." + ) + def test_save_load_optional_components(self): + pass + def test_pipeline_runs(self): pipe = _make_pipeline().to("cpu") @@ -61,8 +135,8 @@ def test_pipeline_runs(self): output_type="seq", ) - self.assertEqual(out.sequences.shape, (2, 24)) - self.assertFalse((out.sequences == 31).any().item()) + assert out.sequences.shape == (2, 24) + assert not (out.sequences == 31).any().item() def test_pipeline_return_tuple(self): pipe = _make_pipeline().to("cpu") @@ -83,8 +157,8 @@ def test_pipeline_return_tuple(self): return_dict=False, ) - self.assertEqual(sequences.shape, (1, 16)) - self.assertIsNone(texts) + assert sequences.shape == (1, 16) + assert texts is None def test_output_type_seq(self): """output_type='seq' should return sequences but no texts.""" @@ -104,9 +178,9 @@ def test_output_type_seq(self): output_type="seq", ) - self.assertIsNotNone(out.sequences) - self.assertEqual(out.sequences.shape, (1, 16)) - self.assertIsNone(out.texts) + assert out.sequences is not None + assert out.sequences.shape == (1, 16) + assert out.texts is None def test_output_type_text_without_tokenizer(self): """output_type='text' without a tokenizer should return texts=None.""" @@ -126,8 +200,8 @@ def test_output_type_text_without_tokenizer(self): output_type="text", ) - self.assertIsNotNone(out.sequences) - self.assertIsNone(out.texts) + assert out.sequences is not None + assert out.texts is None def test_output_type_text_with_tokenizer(self): """output_type='text' with a tokenizer should return decoded texts.""" @@ -155,16 +229,16 @@ def test_output_type_text_with_tokenizer(self): output_type="text", ) - self.assertIsNotNone(out.sequences) - self.assertIsNotNone(out.texts) - self.assertEqual(len(out.texts), 1) - self.assertTrue(out.texts[0].startswith("decoded_")) + assert out.sequences is not None + assert out.texts is not None + assert len(out.texts) == 1 + assert out.texts[0].startswith("decoded_") def test_output_type_invalid_raises(self): """Invalid output_type should raise ValueError.""" pipe = _make_pipeline().to("cpu") - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe( input_ids=torch.tensor([[5, 6, 7, 8]], dtype=torch.long), use_chat_template=False, @@ -186,9 +260,9 @@ def test_prepare_input_ids_from_tensor(self): add_generation_prompt=False, chat_template_kwargs=None, ) - self.assertTrue(torch.equal(result_ids, ids)) - self.assertEqual(result_mask.shape, ids.shape) - self.assertTrue((result_mask == 1).all().item()) + assert torch.equal(result_ids, ids) + assert result_mask.shape == ids.shape + assert (result_mask == 1).all().item() def test_prepare_input_ids_from_1d_tensor(self): pipe = _make_pipeline() @@ -201,12 +275,12 @@ def test_prepare_input_ids_from_1d_tensor(self): add_generation_prompt=False, chat_template_kwargs=None, ) - self.assertEqual(result_ids.shape, (1, 3)) - self.assertEqual(result_mask.shape, (1, 3)) + assert result_ids.shape == (1, 3) + assert result_mask.shape == (1, 3) def test_prepare_input_ids_no_tokenizer_raises(self): pipe = _make_pipeline(tokenizer=None) - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._prepare_input_ids( prompt="hello", messages=None, @@ -220,7 +294,7 @@ def test_prepare_input_ids_both_prompt_and_messages_raises(self): pipe = _make_pipeline() # Manually set tokenizer to a simple object so _prepare_input_ids doesn't short-circuit pipe.tokenizer = type("Tok", (), {"eos_token_id": None, "mask_token_id": None})() - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._prepare_input_ids( prompt="hello", messages=[{"role": "user", "content": "hi"}], @@ -233,7 +307,7 @@ def test_prepare_input_ids_both_prompt_and_messages_raises(self): def test_prepare_input_ids_neither_raises(self): pipe = _make_pipeline() pipe.tokenizer = type("Tok", (), {"eos_token_id": None, "mask_token_id": None})() - with self.assertRaises(ValueError): + with pytest.raises(ValueError): pipe._prepare_input_ids( prompt=None, messages=None, @@ -244,7 +318,7 @@ def test_prepare_input_ids_neither_raises(self): ) -class LLaDA2RegressionTest(unittest.TestCase): +class TestLLaDA2Regression: """Pin the regressions identified in https://github.com/huggingface/diffusers/issues/13598.""" def test_attention_mask_carried_through_for_pre_tokenized_input(self): @@ -278,16 +352,16 @@ def forward(self, input_ids, attention_mask=None, position_ids=None, **kwargs): output_type="seq", ) - self.assertGreater(len(captured), 0) + assert len(captured) > 0 first_mask = captured[0] # Padded prompt positions stay zero in the runtime mask (Issue #1). - self.assertEqual(first_mask[0, 3].item(), 0) - self.assertEqual(first_mask[1, 1].item(), 0) - self.assertEqual(first_mask[1, 2].item(), 0) - self.assertEqual(first_mask[1, 3].item(), 0) + assert first_mask[0, 3].item() == 0 + assert first_mask[1, 1].item() == 0 + assert first_mask[1, 2].item() == 0 + assert first_mask[1, 3].item() == 0 # Real prompt positions stay one. - self.assertEqual(first_mask[0, 0].item(), 1) - self.assertEqual(first_mask[1, 0].item(), 1) + assert first_mask[0, 0].item() == 1 + assert first_mask[1, 0].item() == 1 def test_block_length_routes_into_scheduler_transfer_schedule(self): """Issue #2: the per-call `block_length` must drive the scheduler's `_transfer_schedule`.""" @@ -313,9 +387,9 @@ def cb(pipe, step, timestep, kwargs): callback_on_step_end_tensor_inputs=["transfer_index"], ) # With block_length=num_inference_steps=8 the schedule commits exactly one token per step. - self.assertEqual(commits[0], 1) - self.assertEqual(commits[1], 1) - self.assertEqual(commits[2], 1) + assert commits[0] == 1 + assert commits[1] == 1 + assert commits[2] == 1 def test_callback_tensor_inputs_advertised_keys_resolve(self): """Issue #3: every advertised callback key must be a bound local at callback time.""" @@ -341,7 +415,7 @@ def cb(pipe, step, timestep, kwargs): callback_on_step_end=cb, callback_on_step_end_tensor_inputs=keys, ) - self.assertEqual(set(observed), set(keys)) + assert set(observed) == set(keys) def test_eos_at_first_generated_position_triggers_finished(self): """Issue #4: EOS exactly at index `prompt_length` must mark the row finished.""" @@ -357,7 +431,7 @@ def test_eos_at_first_generated_position_triggers_finished(self): mask_token_id=99, prompt_length=1, ) - self.assertTrue(bool(finished[0].item())) + assert bool(finished[0].item()) def test_finished_rows_are_frozen_for_subsequent_blocks(self): """Issue #5: once a row emits EOS, later blocks must not overwrite its committed tokens.""" @@ -393,7 +467,7 @@ def forward(self, input_ids, attention_mask=None, position_ids=None, **kwargs): output_type="seq", ) # Row 0's first generated tokens must not be overwritten by later-block sampling (token 7). - self.assertNotIn(7, out.sequences[0].tolist()[:2]) + assert 7 not in out.sequences[0].tolist()[:2] def test_progress_bar_disable_is_preserved_after_call(self): """Issue #6: calling the pipeline must not mutate `_progress_bar_config`.""" @@ -412,8 +486,4 @@ def test_progress_bar_disable_is_preserved_after_call(self): eos_early_stop=False, output_type="seq", ) - self.assertEqual(pipe._progress_bar_config, before) - - -if __name__ == "__main__": - unittest.main() + assert pipe._progress_bar_config == before From 125342c72bf75db7278255514579e234f532ea6f Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:50:25 +0200 Subject: [PATCH 6/9] fix the same generator/device bug in add_noise missed this one earlier; two more torch.rand calls with the same CPU-generator-on-CUDA mismatch. --- src/diffusers/schedulers/scheduling_block_refinement.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/diffusers/schedulers/scheduling_block_refinement.py b/src/diffusers/schedulers/scheduling_block_refinement.py index 0bd462d210bc..b7be3483ba7d 100644 --- a/src/diffusers/schedulers/scheduling_block_refinement.py +++ b/src/diffusers/schedulers/scheduling_block_refinement.py @@ -502,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(original_samples, generator) 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] From f3b81f8c3b0fff59598073611b1483d47d9ebb4c Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:50:29 +0200 Subject: [PATCH 7/9] reuse the shared _generator_device helper in DiffusionGemma's pipeline was duplicating the same one-liner inline. --- .../pipelines/diffusion_gemma/pipeline_diffusion_gemma.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py index 99733ccc3c08..a0aff61cba3b 100644 --- a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py +++ b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py @@ -23,6 +23,7 @@ from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...schedulers import BlockRefinementScheduler, DiscreteDDIMScheduler, EntropyBoundScheduler +from ...schedulers.scheduling_utils import _generator_device from ...utils import logging, replace_example_docstring from ..pipeline_utils import DiffusionPipeline from .pipeline_output import DiffusionGemmaPipelineOutput @@ -346,7 +347,7 @@ 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 + rand_device = _generator_device(cur_input_ids, generator) canvas = torch.randint( 0, text_config.vocab_size, (batch_size, canvas_length), device=rand_device, generator=generator ).to(device) From 41418ff27aa8c639613b587dd8bb9f78adfb772a Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 31 Aug 2026 11:50:39 +0200 Subject: [PATCH 8/9] exercise save/load without a tokenizer in the LLaDA2 mixin test was skipped; adapt it to use input_ids instead of prompt so it actually runs. --- tests/pipelines/llada2/test_llada2.py | 36 ++++++++++++++++++++------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/tests/pipelines/llada2/test_llada2.py b/tests/pipelines/llada2/test_llada2.py index c398e1dc7ece..588ad57a09df 100644 --- a/tests/pipelines/llada2/test_llada2.py +++ b/tests/pipelines/llada2/test_llada2.py @@ -4,7 +4,7 @@ from diffusers import BlockRefinementScheduler, LLaDA2Pipeline -from ...testing_utils import enable_full_determinism +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 @@ -72,8 +72,11 @@ def get_dummy_components(self): 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 { - "prompt": "Name a color.", "use_chat_template": False, "generator": self.get_generator(0), "gen_length": self.output_shape[0], @@ -108,13 +111,28 @@ def test_inference_batch_single_identical(self): def test_callback_inputs(self): pass - @pytest.mark.skip( - "Test not supported: drops `tokenizer` and reruns the dummy `prompt` input, but dropping the tokenizer " - "means the caller must switch to pre-tokenized `input_ids` instead (see " - "`test_output_type_text_without_tokenizer` below, which covers this)." - ) - def test_save_load_optional_components(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") From 664f402a2f8c8bda092e6822d0d249959791dda8 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Tue, 1 Sep 2026 08:24:39 +0200 Subject: [PATCH 9/9] fold _generator_device back into each call site per review: it's a one-liner, not worth a shared helper. --- .../diffusion_gemma/pipeline_diffusion_gemma.py | 3 +-- .../schedulers/scheduling_block_refinement.py | 12 ++++++------ .../schedulers/scheduling_discrete_ddim.py | 16 ++++++++-------- .../schedulers/scheduling_entropy_bound.py | 10 +++++----- src/diffusers/schedulers/scheduling_utils.py | 9 --------- 5 files changed, 20 insertions(+), 30 deletions(-) diff --git a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py index a0aff61cba3b..99733ccc3c08 100644 --- a/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py +++ b/src/diffusers/pipelines/diffusion_gemma/pipeline_diffusion_gemma.py @@ -23,7 +23,6 @@ from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...schedulers import BlockRefinementScheduler, DiscreteDDIMScheduler, EntropyBoundScheduler -from ...schedulers.scheduling_utils import _generator_device from ...utils import logging, replace_example_docstring from ..pipeline_utils import DiffusionPipeline from .pipeline_output import DiffusionGemmaPipelineOutput @@ -347,7 +346,7 @@ 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(cur_input_ids, generator) + rand_device = generator.device if generator is not None else device canvas = torch.randint( 0, text_config.vocab_size, (batch_size, canvas_length), device=rand_device, generator=generator ).to(device) diff --git a/src/diffusers/schedulers/scheduling_block_refinement.py b/src/diffusers/schedulers/scheduling_block_refinement.py index b7be3483ba7d..739c9e771098 100644 --- a/src/diffusers/schedulers/scheduling_block_refinement.py +++ b/src/diffusers/schedulers/scheduling_block_refinement.py @@ -20,7 +20,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin, _generator_device +from .scheduling_utils import SchedulerMixin @dataclass @@ -173,9 +173,9 @@ def _sample_from_logits( filtered = BlockRefinementScheduler._top_p_filtering(filtered, top_p=top_p) probs = torch.softmax(filtered.float(), dim=-1) - # `torch.multinomial` requires the generator and the sampled tensor's device to match; `_generator_device` - # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. - rand_device = _generator_device(probs, 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) @@ -288,7 +288,7 @@ def step( prev_sample = torch.where(transfer_index | editing_transfer_index, sampled_tokens, sample) self._committed = committed | transfer_index - rand_device = _generator_device(sample, generator) + 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=rand_device, generator=generator ).to(sample.device) @@ -502,7 +502,7 @@ def add_noise( masked_rev = torch.zeros_like(original_samples, dtype=torch.bool) valid = attention_mask.to(dtype=torch.bool) - rand_device = _generator_device(original_samples, generator) + 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 diff --git a/src/diffusers/schedulers/scheduling_discrete_ddim.py b/src/diffusers/schedulers/scheduling_discrete_ddim.py index 96f4d22cb21b..3fa598c93321 100644 --- a/src/diffusers/schedulers/scheduling_discrete_ddim.py +++ b/src/diffusers/schedulers/scheduling_discrete_ddim.py @@ -21,7 +21,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin, _generator_device +from .scheduling_utils import SchedulerMixin @dataclass @@ -117,9 +117,9 @@ def _sample_from_logits( token = flat_logits.argmax(dim=-1, keepdim=True) else: scaled_probs = torch.softmax(flat_logits.float() / temperature, dim=-1) - # `torch.multinomial` requires the generator and the sampled tensor's device to match; `_generator_device` - # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. - rand_device = _generator_device(scaled_probs, 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 ) @@ -205,7 +205,7 @@ 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) - rand_device = _generator_device(route_probs, generator) + 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) @@ -236,7 +236,7 @@ def _select_positions( k_eff = min(max(1, int(self.config.corrector_k)), seq_len) if selection == "random": - rand_device = _generator_device(sample, 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 @@ -252,7 +252,7 @@ def _select_positions( raise ValueError(f"Unknown `corrector_selection`: {selection!r}.") keys = confidence / float(self.config.corrector_selection_tau) - rand_device = _generator_device(keys, generator) + 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 @@ -307,7 +307,7 @@ 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() - rand_device = _generator_device(chosen_probs, generator) + 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) diff --git a/src/diffusers/schedulers/scheduling_entropy_bound.py b/src/diffusers/schedulers/scheduling_entropy_bound.py index 4c39263ba518..f59c81710b12 100644 --- a/src/diffusers/schedulers/scheduling_entropy_bound.py +++ b/src/diffusers/schedulers/scheduling_entropy_bound.py @@ -20,7 +20,7 @@ from ..configuration_utils import ConfigMixin, register_to_config from ..utils import BaseOutput -from .scheduling_utils import SchedulerMixin, _generator_device +from .scheduling_utils import SchedulerMixin @dataclass @@ -110,9 +110,9 @@ def _sample_from_logits( token = flat_logits.argmax(dim=-1, keepdim=True) else: scaled_probs = torch.softmax(flat_logits.float() / temperature, dim=-1) - # `torch.multinomial` requires the generator and the sampled tensor's device to match; `_generator_device` - # gives the CPU-generator portability `randn_tensor` gives the continuous samplers. - rand_device = _generator_device(scaled_probs, 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 ) @@ -171,7 +171,7 @@ def step( input=torch.zeros_like(sorted_accepted), dim=-1, index=sorted_indices, src=sorted_accepted ) - rand_device = _generator_device(sample, generator) + 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=rand_device, generator=generator ).to(sample.device) diff --git a/src/diffusers/schedulers/scheduling_utils.py b/src/diffusers/schedulers/scheduling_utils.py index 1d191a2dbbfc..8cdb21c2f011 100644 --- a/src/diffusers/schedulers/scheduling_utils.py +++ b/src/diffusers/schedulers/scheduling_utils.py @@ -27,15 +27,6 @@ SCHEDULER_CONFIG_NAME = "scheduler_config.json" -def _generator_device(tensor: torch.Tensor, generator: torch.Generator | None) -> torch.device: - """Device to sample on for a discrete (`torch.randint`/`torch.multinomial`/`torch.rand`) draw: the generator's - own device when one is passed, else `tensor`'s device. Mirrors `randn_tensor`'s CPU-generator portability (a CPU - generator, the recommended default for reproducible pipeline calls, samples on CPU regardless of `tensor`'s device) - for the discrete samplers, which have no `randn_tensor` equivalent to call into. - """ - return generator.device if generator is not None else tensor.device - - # NOTE: We make this type an enum because it simplifies usage in docs and prevents # circular imports when used for `_compatibles` within the schedulers module. # When it's used as a type in pipelines, it really is a Union because the actual