nucleusmoe_image model/pipeline review

[hlky]

nucleusmoe_image model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Issue 1: Batched prompts can crash CFG when negative_prompt is a string

Affected code:

diffusers/src/diffusers/pipelines/nucleusmoe_image/pipeline_nucleusmoe_image.py

Lines 492 to 518 in 0f1abc4

	has_neg_prompt = negative_prompt is not None or (
	negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
	)
	do_cfg = guidance_scale > 1
	if do_cfg and not has_neg_prompt:
	negative_prompt = [""] * batch_size
	prompt_embeds, prompt_embeds_mask = self.encode_prompt(
	prompt=prompt,
	prompt_embeds=prompt_embeds,
	prompt_embeds_mask=prompt_embeds_mask,
	device=device,
	num_images_per_prompt=num_images_per_prompt,
	max_sequence_length=max_sequence_length,
	return_index=return_index,
	)
	if do_cfg:
	negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
	prompt=negative_prompt,
	prompt_embeds=negative_prompt_embeds,
	prompt_embeds_mask=negative_prompt_embeds_mask,
	device=device,
	num_images_per_prompt=num_images_per_prompt,
	max_sequence_length=max_sequence_length,
	return_index=return_index,
	)

diffusers/src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Lines 317 to 318 in 0f1abc4

	joint_key = torch.cat([img_key, txt_key], dim=1)
	joint_value = torch.cat([img_value, txt_value], dim=1)

Problem:
For prompt=["a", "b"], batch_size is 2. If the caller passes negative_prompt="bad" with guidance_scale > 1, the negative prompt is encoded as batch size 1, then passed to the transformer with latents of batch size 2. The attention processor later concatenates image and text K/V tensors and raises a batch-size mismatch.

Impact:
Common batched text-to-image usage crashes during CFG instead of either broadcasting the scalar negative prompt or raising a clear validation error.

Reproduction:

import torch
from diffusers import NucleusMoEImageTransformer2DModel

model = NucleusMoEImageTransformer2DModel(
 patch_size=2, in_channels=16, out_channels=4, num_layers=1,
 attention_head_dim=16, num_attention_heads=4, joint_attention_dim=16,
 axes_dims_rope=(8, 4, 4), moe_enabled=False, capacity_factors=[8.0],
).eval()

latents = torch.randn(2, 16, 16) # prompt=["a", "b"]
negative_embeds = torch.randn(1, 8, 16) # negative_prompt="bad"

with torch.no_grad():
 model(
 hidden_states=latents,
 timestep=torch.ones(2),
 encoder_hidden_states=negative_embeds,
 encoder_hidden_states_mask=torch.ones(1, 8, dtype=torch.long),
 img_shapes=[(1, 4, 4), (1, 4, 4)],
 )

Relevant precedent:

diffusers/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

Lines 744 to 748 in 0f1abc4

	if self.do_classifier_free_guidance:
	negative_prompt = ""
	if prompt is not None and isinstance(prompt, list):
	negative_prompt = [negative_prompt] * len(prompt)
	negative_prompt_embeds, negative_text_ids = self.encode_prompt(

diffusers/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

Lines 455 to 472 in 0f1abc4

	uncond_tokens: list[str]
	if negative_prompt is None:
	uncond_tokens = [""] * batch_size
	elif prompt is not None and type(prompt) is not type(negative_prompt):
	raise TypeError(
	f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
	f" {type(prompt)}."
	)
	elif isinstance(negative_prompt, str):
	uncond_tokens = [negative_prompt]
	elif batch_size != len(negative_prompt):
	raise ValueError(
	f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
	f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
	" the batch size of `prompt`."
	)
	else:
	uncond_tokens = negative_prompt

Suggested fix:

if do_cfg and negative_prompt_embeds is None:
 if negative_prompt is None:
 negative_prompt = [""] * batch_size
 elif isinstance(negative_prompt, str) and batch_size > 1:
 negative_prompt = [negative_prompt] * batch_size
 elif isinstance(negative_prompt, list) and len(negative_prompt) != batch_size:
 raise ValueError(
 f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size {batch_size}."
 )

Issue 2: return_index rejects valid hidden-state indices and ignores 0

Affected code:

diffusers/src/diffusers/pipelines/nucleusmoe_image/pipeline_nucleusmoe_image.py

Lines 214 to 235 in 0f1abc4

	device = device or self._execution_device
	return_index = return_index or self.default_return_index
	if prompt_embeds is None:
	prompt = [prompt] if isinstance(prompt, str) else prompt
	formatted = [self._format_prompt(p) for p in prompt]
	inputs = self.processor(
	text=formatted,
	padding="longest",
	pad_to_multiple_of=8,
	max_length=max_sequence_length,
	truncation=True,
	return_attention_mask=True,
	return_tensors="pt",
	).to(device=device)
	prompt_embeds_mask = inputs.attention_mask
	outputs = self.text_encoder(**inputs, use_cache=False, return_dict=True, output_hidden_states=True)
	prompt_embeds = outputs.hidden_states[return_index]
	prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)

diffusers/src/diffusers/pipelines/nucleusmoe_image/pipeline_nucleusmoe_image.py

Lines 296 to 300 in 0f1abc4

	if return_index is not None and abs(return_index) >= self.text_encoder.config.text_config.num_hidden_layers:
	raise ValueError(
	f"absolute value of `return_index` cannot be >= {self.text_encoder.config.text_config.num_hidden_layers} "
	f"but is {abs(return_index)}"
	)

Problem:
return_index = return_index or self.default_return_index makes explicit return_index=0 impossible. Validation also uses abs(return_index) >= num_hidden_layers, but Transformer hidden states include the embeddings plus one entry per layer, so valid indices include num_hidden_layers and -(num_hidden_layers + 1) through -1.

Impact:
Users cannot select hidden state 0, and valid negative layer selections can be rejected before encoding.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import NucleusMoEImagePipeline

class Batch(dict):
 @property
 def attention_mask(self):
 return self["attention_mask"]
 def to(self, device=None):
 return self

class FakeProcessor:
 def apply_chat_template(self, *args, **kwargs):
 return "formatted"
 def __call__(self, **kwargs):
 return Batch(input_ids=torch.tensor([[1]]), attention_mask=torch.tensor([[1]]))

class FakeTextEncoder:
 dtype = torch.float32
 config = SimpleNamespace(text_config=SimpleNamespace(num_hidden_layers=8))
 def __call__(self, **kwargs):
 return SimpleNamespace(hidden_states=[torch.full((1, 1, 1), i, dtype=torch.float32) for i in range(9)])

pipe = object.__new__(NucleusMoEImagePipeline)
pipe.processor = FakeProcessor()
pipe.text_encoder = FakeTextEncoder()
pipe.default_return_index = -8

embeds, _ = pipe.encode_prompt("x", device=torch.device("cpu"), return_index=0)
print(embeds.item()) # 1.0, expected 0.0

pipe.vae_scale_factor = 8
pipe._callback_tensor_inputs = ["latents", "prompt_embeds"]
pipe.check_inputs("x", 64, 64, return_index=-8) # raises, but hidden_states[-8] is valid for length 9

Suggested fix:

return_index = self.default_return_index if return_index is None else return_index

num_hidden_states = self.text_encoder.config.text_config.num_hidden_layers + 1
if return_index is not None and not (-num_hidden_states <= return_index < num_hidden_states):
 raise ValueError(
 f"`return_index` must be in [{-num_hidden_states}, {num_hidden_states - 1}], but is {return_index}."
 )

Issue 3: Precomputed prompt embeddings ignore max_sequence_length

Affected code:

diffusers/src/diffusers/pipelines/nucleusmoe_image/pipeline_nucleusmoe_image.py

Lines 236 to 249 in 0f1abc4

	else:
	prompt_embeds = prompt_embeds.to(device=device)
	if prompt_embeds_mask is not None:
	prompt_embeds_mask = prompt_embeds_mask.to(device=device)
	if num_images_per_prompt > 1:
	prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0)
	if prompt_embeds_mask is not None:
	prompt_embeds_mask = prompt_embeds_mask.repeat_interleave(num_images_per_prompt, dim=0)
	if prompt_embeds_mask is not None and prompt_embeds_mask.all():
	prompt_embeds_mask = None
	return prompt_embeds, prompt_embeds_mask

Problem:
When prompt_embeds are supplied directly, encode_prompt moves them to the device but never slices prompt_embeds or prompt_embeds_mask to max_sequence_length. Freshly encoded prompts are truncated by the processor, so the two code paths disagree.

Impact:
Calls that use cached/precomputed embeddings can feed longer text context than requested, changing RoPE text positions, attention length, memory use, and output parity with the normal prompt path.

Reproduction:

import torch
from diffusers import NucleusMoEImagePipeline

pipe = object.__new__(NucleusMoEImagePipeline)
pipe.default_return_index = -8

embeds, mask = pipe.encode_prompt(
 prompt_embeds=torch.randn(1, 8, 4),
 prompt_embeds_mask=torch.tensor([[1, 1, 1, 1, 0, 0, 0, 0]]),
 device=torch.device("cpu"),
 max_sequence_length=4,
)
print(embeds.shape, mask.shape) # torch.Size([1, 8, 4]) torch.Size([1, 8])

Relevant precedent:

diffusers/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py

Lines 256 to 264 in 0f1abc4

	prompt_embeds = prompt_embeds[:, :max_sequence_length]
	_, seq_len, _ = prompt_embeds.shape
	prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
	prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
	if prompt_embeds_mask is not None:
	prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
	prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
	prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)

Suggested fix:

if max_sequence_length is None:
 max_sequence_length = self.default_max_sequence_length

prompt_embeds = prompt_embeds[:, :max_sequence_length]
if prompt_embeds_mask is not None:
 prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]

Issue 4: Generic Attention APIs can replace the custom processor and break NucleusMoE attention

Affected code:

diffusers/src/diffusers/models/transformers/transformer_nucleusmoe_image.py

Lines 643 to 657 in 0f1abc4

	self.attn = Attention(
	query_dim=dim,
	heads=num_attention_heads,
	kv_heads=num_key_value_heads,
	dim_head=attention_head_dim,
	added_kv_proj_dim=dim,
	added_proj_bias=False,
	out_dim=dim,
	out_bias=False,
	bias=False,
	processor=NucleusMoEAttnProcessor2_0(),
	qk_norm=qk_norm,
	eps=eps,
	context_pre_only=None,
	)

Problem:
The block instantiates the generic Attention class with NucleusMoEAttnProcessor2_0. Generic Attention exposes processor-changing APIs such as set_attention_slice, which replace the Nucleus processor with incompatible generic processors. The review rules require a model-local attention class using AttentionModuleMixin, _default_processor_cls, and _available_processors to avoid exactly this class of breakage.

Impact:
Lower-level attention APIs can leave the model in a broken state. The pipeline fast test for attention slicing does not catch this because DiffusionPipeline.enable_attention_slicing() only calls components that expose set_attention_slice; NucleusMoEImageTransformer2DModel does not, so the pipeline-level call is effectively a no-op.

Reproduction:

import torch
from diffusers import NucleusMoEImageTransformer2DModel

model = NucleusMoEImageTransformer2DModel(
 patch_size=2, in_channels=16, out_channels=4, num_layers=1,
 attention_head_dim=16, num_attention_heads=4, joint_attention_dim=16,
 axes_dims_rope=(8, 4, 4), moe_enabled=False, capacity_factors=[8.0],
).eval()

model.transformer_blocks[0].attn.set_attention_slice(1)

with torch.no_grad():
 model(
 hidden_states=torch.randn(1, 16, 16),
 timestep=torch.ones(1),
 encoder_hidden_states=torch.randn(1, 8, 16),
 encoder_hidden_states_mask=torch.ones(1, 8, dtype=torch.long),
 img_shapes=[(1, 4, 4)],
 )

Relevant precedent:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 493 to 548 in 0f1abc4

	class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
	_default_processor_cls = Flux2AttnProcessor
	_available_processors = [Flux2AttnProcessor, Flux2KVAttnProcessor]
	def __init__(
	self,
	query_dim: int,
	heads: int = 8,
	dim_head: int = 64,
	dropout: float = 0.0,
	bias: bool = False,
	added_kv_proj_dim: int | None = None,
	added_proj_bias: bool | None = True,
	out_bias: bool = True,
	eps: float = 1e-5,
	out_dim: int = None,
	elementwise_affine: bool = True,
	processor=None,
	):
	super().__init__()
	self.head_dim = dim_head
	self.inner_dim = out_dim if out_dim is not None else dim_head * heads
	self.query_dim = query_dim
	self.out_dim = out_dim if out_dim is not None else query_dim
	self.heads = out_dim // dim_head if out_dim is not None else heads
	self.use_bias = bias
	self.dropout = dropout
	self.added_kv_proj_dim = added_kv_proj_dim
	self.added_proj_bias = added_proj_bias
	self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
	self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
	self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
	# QK Norm
	self.norm_q = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
	self.norm_k = torch.nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
	self.to_out = torch.nn.ModuleList([])
	self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
	self.to_out.append(torch.nn.Dropout(dropout))
	if added_kv_proj_dim is not None:
	self.norm_added_q = torch.nn.RMSNorm(dim_head, eps=eps)
	self.norm_added_k = torch.nn.RMSNorm(dim_head, eps=eps)
	self.add_q_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
	self.add_k_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
	self.add_v_proj = torch.nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
	self.to_add_out = torch.nn.Linear(self.inner_dim, query_dim, bias=out_bias)
	if processor is None:
	processor = self._default_processor_cls()
	self.set_processor(processor)

diffusers/src/diffusers/models/transformers/transformer_wan.py

Lines 175 to 177 in 0f1abc4

	class WanAttention(torch.nn.Module, AttentionModuleMixin):
	_default_processor_cls = WanAttnProcessor
	_available_processors = [WanAttnProcessor]

Suggested fix:
Implement a NucleusMoEAttention(torch.nn.Module, AttentionModuleMixin) in transformer_nucleusmoe_image.py, move the projections/norms used by NucleusMoEAttnProcessor2_0 into it, set _default_processor_cls = NucleusMoEAttnProcessor2_0, and restrict _available_processors to compatible processors.

Issue 5: Slow tests are missing

Affected code:

diffusers/tests/pipelines/nucleusmoe_image/test_nucleusmoe_image.py

Lines 38 to 337 in 0f1abc4

	class NucleusMoEImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
	pipeline_class = NucleusMoEImagePipeline
	params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
	batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
	image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
	image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
	required_optional_params = frozenset(
	[
	"num_inference_steps",
	"generator",
	"latents",
	"return_dict",
	"callback_on_step_end",
	"callback_on_step_end_tensor_inputs",
	]
	)
	supports_dduf = False
	test_xformers_attention = False
	test_layerwise_casting = True
	test_group_offloading = True
	def get_dummy_components(self):
	torch.manual_seed(0)
	transformer = NucleusMoEImageTransformer2DModel(
	patch_size=2,
	in_channels=16,
	out_channels=4,
	num_layers=2,
	attention_head_dim=16,
	num_attention_heads=4,
	joint_attention_dim=16,
	axes_dims_rope=(8, 4, 4),
	moe_enabled=False,
	capacity_factors=[8.0, 8.0],
	)
	torch.manual_seed(0)
	z_dim = 4
	vae = AutoencoderKLQwenImage(
	base_dim=z_dim * 6,
	z_dim=z_dim,
	dim_mult=[1, 2, 4],
	num_res_blocks=1,
	temperal_downsample=[False, True],
	# fmt: off
	latents_mean=[0.0] * z_dim,
	latents_std=[1.0] * z_dim,
	# fmt: on
	)
	torch.manual_seed(0)
	scheduler = FlowMatchEulerDiscreteScheduler()
	torch.manual_seed(0)
	config = Qwen3VLConfig(
	text_config={
	"hidden_size": 16,
	"intermediate_size": 16,
	"num_hidden_layers": 8,
	"num_attention_heads": 2,
	"num_key_value_heads": 2,
	"rope_scaling": {
	"mrope_section": [1, 1, 2],
	"rope_type": "default",
	"type": "default",
	},
	"rope_theta": 1000000.0,
	"vocab_size": 151936,
	"head_dim": 8,
	},
	vision_config={
	"depth": 2,
	"hidden_size": 16,
	"intermediate_size": 16,
	"num_heads": 2,
	"out_channels": 16,
	},
	)
	text_encoder = Qwen3VLForConditionalGeneration(config).eval()
	processor = Qwen3VLProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
	components = {
	"transformer": transformer,
	"vae": vae,
	"scheduler": scheduler,
	"text_encoder": text_encoder,
	"processor": processor,
	}
	return components
	def get_dummy_inputs(self, device, seed=0):
	if str(device).startswith("mps"):
	generator = torch.manual_seed(seed)
	else:
	generator = torch.Generator(device=device).manual_seed(seed)
	inputs = {
	"prompt": "A cat sitting on a mat",
	"negative_prompt": "bad quality",
	"generator": generator,
	"num_inference_steps": 2,
	"return_index": -1,
	"guidance_scale": 1.0,
	"height": 32,
	"width": 32,
	"max_sequence_length": 16,
	"output_type": "pt",
	}
	return inputs
	def test_inference(self):
	device = "cpu"
	components = self.get_dummy_components()
	pipe = self.pipeline_class(**components)
	pipe.to(device)
	pipe.set_progress_bar_config(disable=None)
	inputs = self.get_dummy_inputs(device)
	image = pipe(**inputs).images
	generated_image = image[0]
	self.assertEqual(generated_image.shape, (3, 32, 32))
	def test_inference_batch_single_identical(self):
	self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-1)
	def test_true_cfg(self):
	device = "cpu"
	components = self.get_dummy_components()
	pipe = self.pipeline_class(**components)
	pipe.to(device)
	pipe.set_progress_bar_config(disable=None)
	inputs = self.get_dummy_inputs(device)
	inputs["guidance_scale"] = 4.0
	inputs["negative_prompt"] = "low quality"
	image = pipe(**inputs).images
	self.assertEqual(image[0].shape, (3, 32, 32))
	def test_prompt_embeds(self):
	device = "cpu"
	components = self.get_dummy_components()
	pipe = self.pipeline_class(**components)
	pipe.to(device)
	pipe.set_progress_bar_config(disable=None)
	inputs = self.get_dummy_inputs(device)
	prompt_embeds, prompt_embeds_mask = pipe.encode_prompt(
	prompt=inputs["prompt"],
	device=device,
	max_sequence_length=inputs["max_sequence_length"],
	)
	inputs_with_embeds = self.get_dummy_inputs(device)
	inputs_with_embeds.pop("prompt")
	inputs_with_embeds["prompt_embeds"] = prompt_embeds
	inputs_with_embeds["prompt_embeds_mask"] = prompt_embeds_mask
	image = pipe(**inputs_with_embeds).images
	self.assertEqual(image[0].shape, (3, 32, 32))
	def test_attention_slicing_forward_pass(
	self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
	):
	# PipelineTesterMixin compares outputs with assert_mean_pixel_difference, which assumes HWC numpy/PIL layout.
	# With output_type="pt", tensors are CHW; numpy_to_pil then fails. Match QwenImage: only assert max diff.
	if not self.test_attention_slicing:
	return
	components = self.get_dummy_components()
	pipe = self.pipeline_class(**components)
	for component in pipe.components.values():
	if hasattr(component, "set_default_attn_processor"):
	component.set_default_attn_processor()
	pipe.to(torch_device)
	pipe.set_progress_bar_config(disable=None)
	generator_device = "cpu"
	inputs = self.get_dummy_inputs(generator_device)
	output_without_slicing = pipe(**inputs)[0]
	pipe.enable_attention_slicing(slice_size=1)
	inputs = self.get_dummy_inputs(generator_device)
	output_with_slicing1 = pipe(**inputs)[0]
	pipe.enable_attention_slicing(slice_size=2)
	inputs = self.get_dummy_inputs(generator_device)
	output_with_slicing2 = pipe(**inputs)[0]
	if test_max_difference:
	max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
	max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
	self.assertLess(
	max(max_diff1, max_diff2),
	expected_max_diff,
	"Attention slicing should not affect the inference results",
	)
	def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4):
	# PipelineTesterMixin only keeps components whose keys contain "text" or "tokenizer"; this pipeline also
	# needs `processor` for encode_prompt (apply_chat_template). Mirror the mixin with that key included.
	if not hasattr(self.pipeline_class, "encode_prompt"):
	return
	components = self.get_dummy_components()
	for key in components:
	if "text_encoder" in key and hasattr(components[key], "eval"):
	components[key].eval()
	def _is_text_stack_component(k):
	return "text" in k or "tokenizer" in k or k == "processor"
	components_with_text_encoders = {}
	for k in components:
	if _is_text_stack_component(k):
	components_with_text_encoders[k] = components[k]
	else:
	components_with_text_encoders[k] = None
	pipe_with_just_text_encoder = self.pipeline_class(**components_with_text_encoders)
	pipe_with_just_text_encoder = pipe_with_just_text_encoder.to(torch_device)
	inputs = self.get_dummy_inputs(torch_device)
	encode_prompt_signature = inspect.signature(pipe_with_just_text_encoder.encode_prompt)
	encode_prompt_parameters = list(encode_prompt_signature.parameters.values())
	required_params = []
	for param in encode_prompt_parameters:
	if param.name == "self" or param.name == "kwargs":
	continue
	if param.default is inspect.Parameter.empty:
	required_params.append(param.name)
	encode_prompt_param_names = [p.name for p in encode_prompt_parameters if p.name != "self"]
	input_keys = list(inputs.keys())
	encode_prompt_inputs = {k: inputs.pop(k) for k in input_keys if k in encode_prompt_param_names}
	pipe_call_signature = inspect.signature(pipe_with_just_text_encoder.__call__)
	pipe_call_parameters = pipe_call_signature.parameters
	for required_param_name in required_params:
	if required_param_name not in encode_prompt_inputs:
	pipe_call_param = pipe_call_parameters.get(required_param_name, None)
	if pipe_call_param is not None and pipe_call_param.default is not inspect.Parameter.empty:
	encode_prompt_inputs[required_param_name] = pipe_call_param.default
	elif extra_required_param_value_dict is not None and isinstance(extra_required_param_value_dict, dict):
	encode_prompt_inputs[required_param_name] = extra_required_param_value_dict[required_param_name]
	else:
	raise ValueError(
	f"Required parameter '{required_param_name}' in "
	f"encode_prompt has no default in either encode_prompt or __call__."
	)
	with torch.no_grad():
	encoded_prompt_outputs = pipe_with_just_text_encoder.encode_prompt(**encode_prompt_inputs)
	ast_visitor = ReturnNameVisitor()
	encode_prompt_tree = ast_visitor.get_ast_tree(cls=self.pipeline_class)
	ast_visitor.visit(encode_prompt_tree)
	prompt_embed_kwargs = ast_visitor.return_names
	prompt_embeds_kwargs = dict(zip(prompt_embed_kwargs, encoded_prompt_outputs))
	adapted_prompt_embeds_kwargs = {
	k: prompt_embeds_kwargs.pop(k) for k in list(prompt_embeds_kwargs.keys()) if k in pipe_call_parameters
	}
	components_with_text_encoders = {}
	for k in components:
	if _is_text_stack_component(k):
	components_with_text_encoders[k] = None
	else:
	components_with_text_encoders[k] = components[k]
	pipe_without_text_encoders = self.pipeline_class(**components_with_text_encoders).to(torch_device)
	pipe_without_tes_inputs = {**inputs, **adapted_prompt_embeds_kwargs}
	if (
	pipe_call_parameters.get("negative_prompt", None) is not None
	and pipe_call_parameters.get("negative_prompt").default is not None
	):
	pipe_without_tes_inputs.update({"negative_prompt": None})
	if (
	pipe_call_parameters.get("prompt", None) is not None
	and pipe_call_parameters.get("prompt").default is inspect.Parameter.empty
	and pipe_call_parameters.get("prompt_embeds", None) is not None
	and pipe_call_parameters.get("prompt_embeds").default is None
	):
	pipe_without_tes_inputs.update({"prompt": None})
	pipe_out = pipe_without_text_encoders(**pipe_without_tes_inputs)[0]
	full_pipe = self.pipeline_class(**components).to(torch_device)
	inputs = self.get_dummy_inputs(torch_device)
	pipe_out_2 = full_pipe(**inputs)[0]
	if isinstance(pipe_out, np.ndarray) and isinstance(pipe_out_2, np.ndarray):
	self.assertTrue(np.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))
	elif isinstance(pipe_out, torch.Tensor) and isinstance(pipe_out_2, torch.Tensor):
	self.assertTrue(torch.allclose(pipe_out, pipe_out_2, atol=atol, rtol=rtol))

diffusers/tests/models/transformers/test_models_transformer_nucleusmoe_image.py

Lines 104 to 220 in 0f1abc4

	class TestNucleusMoEImageTransformer(NucleusMoEImageTransformerTesterConfig, ModelTesterMixin):
	def test_with_attention_mask(self):
	init_dict = self.get_init_dict()
	inputs = self.get_dummy_inputs()
	model = self.model_class(**init_dict).to(torch_device)
	# Mask out some text tokens
	mask = inputs["encoder_hidden_states_mask"].clone()
	mask[:, 4:] = 0
	inputs["encoder_hidden_states_mask"] = mask
	with torch.no_grad():
	output = model(**inputs)
	assert output.sample.shape[1] == inputs["hidden_states"].shape[1]
	def test_without_attention_mask(self):
	init_dict = self.get_init_dict()
	inputs = self.get_dummy_inputs()
	model = self.model_class(**init_dict).to(torch_device)
	inputs["encoder_hidden_states_mask"] = None
	with torch.no_grad():
	output = model(**inputs)
	assert output.sample.shape[1] == inputs["hidden_states"].shape[1]
	class TestNucleusMoEImageTransformerMemory(NucleusMoEImageTransformerTesterConfig, MemoryTesterMixin):
	"""Memory optimization tests for NucleusMoE Image Transformer."""
	class TestNucleusMoEImageTransformerTraining(NucleusMoEImageTransformerTesterConfig, TrainingTesterMixin):
	"""Training tests for NucleusMoE Image Transformer."""
	class TestNucleusMoEImageTransformerAttention(NucleusMoEImageTransformerTesterConfig, AttentionTesterMixin):
	"""Attention processor tests for NucleusMoE Image Transformer."""
	class TestNucleusMoEImageTransformerLoRA(NucleusMoEImageTransformerTesterConfig, LoraTesterMixin):
	"""LoRA adapter tests for NucleusMoE Image Transformer."""
	class TestNucleusMoEImageTransformerLoRAHotSwap(
	NucleusMoEImageTransformerTesterConfig, LoraHotSwappingForModelTesterMixin
	):
	"""LoRA hot-swapping tests for NucleusMoE Image Transformer."""
	@property
	def different_shapes_for_compilation(self):
	return [(4, 4), (4, 8), (8, 8)]
	def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict:
	batch_size = 1
	in_channels = 16
	joint_attention_dim = 16
	sequence_length = 8
	hidden_states = randn_tensor(
	(batch_size, height * width, in_channels), generator=self.generator, device=torch_device
	)
	encoder_hidden_states = randn_tensor(
	(batch_size, sequence_length, joint_attention_dim), generator=self.generator, device=torch_device
	)
	encoder_hidden_states_mask = torch.ones((batch_size, sequence_length), dtype=torch.long).to(torch_device)
	timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
	img_shapes = [(1, height, width)] * batch_size
	return {
	"hidden_states": hidden_states,
	"encoder_hidden_states": encoder_hidden_states,
	"encoder_hidden_states_mask": encoder_hidden_states_mask,
	"timestep": timestep,
	"img_shapes": img_shapes,
	}
	class TestNucleusMoEImageTransformerCompile(NucleusMoEImageTransformerTesterConfig, TorchCompileTesterMixin):
	"""Torch compile tests for NucleusMoE Image Transformer."""
	@property
	def different_shapes_for_compilation(self):
	return [(4, 4), (4, 8), (8, 8)]
	def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict:
	batch_size = 1
	in_channels = 16
	joint_attention_dim = 16
	sequence_length = 8
	hidden_states = randn_tensor(
	(batch_size, height * width, in_channels), generator=self.generator, device=torch_device
	)
	encoder_hidden_states = randn_tensor(
	(batch_size, sequence_length, joint_attention_dim), generator=self.generator, device=torch_device
	)
	encoder_hidden_states_mask = torch.ones((batch_size, sequence_length), dtype=torch.long).to(torch_device)
	timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size)
	img_shapes = [(1, height, width)] * batch_size
	return {
	"hidden_states": hidden_states,
	"encoder_hidden_states": encoder_hidden_states,
	"encoder_hidden_states_mask": encoder_hidden_states_mask,
	"timestep": timestep,
	"img_shapes": img_shapes,
	}
	class TestNucleusMoEImageTransformerBitsAndBytes(NucleusMoEImageTransformerTesterConfig, BitsAndBytesTesterMixin):
	"""BitsAndBytes quantization tests for NucleusMoE Image Transformer."""
	class TestNucleusMoEImageTransformerTorchAo(NucleusMoEImageTransformerTesterConfig, TorchAoTesterMixin):
	"""TorchAO quantization tests for NucleusMoE Image Transformer."""

Problem:
The target family has fast pipeline and model tests, but no @slow coverage for the real published checkpoint.

Impact:
The integration does not verify real-component loading, real tokenizer/processor behavior, checkpoint config compatibility, or an expected output slice. This leaves conversion/parity regressions undetected.

Reproduction:

from pathlib import Path

for path in sorted(Path("tests").rglob("*nucleusmoe*")):
 if path.is_file() and "__pycache__" not in path.parts:
 text = path.read_text(encoding="utf-8", errors="ignore")
 print(path, "@slow" in text)

Suggested fix:
Add a slow pipeline test that loads NucleusAI/NucleusMoE-Image with the documented dtype/device path, runs a small deterministic prompt, and checks an output slice or image statistics. If full generation is too expensive, add a slow load/encode/one-step smoke test with explicit justification.

Issue 6: Public pipeline has no docs or examples

Affected code:

diffusers/src/diffusers/pipelines/auto_pipeline.py

Line 80 in 0f1abc4

from .nucleusmoe_image import NucleusMoEImagePipeline

diffusers/src/diffusers/pipelines/auto_pipeline.py

Line 183 in 0f1abc4

("nucleusmoe-image", NucleusMoEImagePipeline),

Problem:
NucleusMoEImagePipeline is exported and registered in auto-pipeline mappings, but there are no matching docs or examples under docs/ or examples/.

Impact:
The public pipeline lacks discoverable API documentation, output docs, usage notes for Qwen3-VL processor requirements, dtype/offload recommendations, and any documented limitations.

Reproduction:

from pathlib import Path

patterns = ("NucleusMoE", "nucleusmoe", "NucleusMoEImage")
for root in [Path("docs"), Path("examples")]:
 hits = []
 for path in root.rglob("*"):
 if path.is_file() and path.suffix.lower() in {".md", ".mdx", ".py", ".rst"}:
 text = path.read_text(encoding="utf-8", errors="ignore")
 if any(pattern in text for pattern in patterns):
 hits.append(str(path))
 print(root, hits)

Suggested fix:
Add a docs page, include it in the pipeline docs toctree, document NucleusMoEImagePipeline, NucleusMoEImagePipelineOutput, usage with NucleusAI/NucleusMoE-Image, and recommended memory/offload settings.

Duplicate-search status

Searched GitHub Issues and PRs in huggingface/diffusers for NucleusMoEImage, NucleusMoEImagePipeline, NucleusMoEImageTransformer2DModel, nucleusmoe_image, pipeline_nucleusmoe_image, return_index, negative_prompt, prompt_embeds max_sequence_length, slow tests, and docs. No duplicate issue/PR was found for these findings. The only related PR surfaced was the merged integration PR: #13317.

Verification notes: top-level imports work in the current .venv, and a tiny transformer save/load round trip produced max diff 0.0. Targeted pytest collection failed in this .venv before running tests because the installed Windows PyTorch build lacks torch._C._distributed_c10d.