Run DiffusionGemma and LLaDA2 pipeline tests through the shared PipelineTesterMixin

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

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",

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

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]
 

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

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:

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