Skip to content

cogview3 model/pipeline review #13648

Description

@hlky

cogview3 model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search: checked GitHub Issues/PRs for cogview3, affected class/output names, prompt-embed batch failures, size validation, attention backend behavior, pooled_projection_dim, stale checkpoint IDs, and fp16 black-image failures. Related but not full duplicates: PR #10211 fixed only the pipeline example checkpoint ID; issue #10343 covers CogView3 fp16 black images.

Local execution note: targeted reproductions were run with .venv. Full fast test collection did not reach target code because this environment's Torch build lacks torch._C._distributed_c10d.

Issue 1: Lazy export points to a non-existent output class

Affected code:

_import_structure= {"pipeline_output": ["CogView3PlusPipelineOutput"]}

@dataclass
classCogView3PipelineOutput(BaseOutput):
"""
Output class for CogView3 pipelines.
Args:
images (`list[PIL.Image.Image]` or `np.ndarray`)
list of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
"""
images: list[PIL.Image.Image] |np.ndarray

Problem:
cogview3.__init__ exports CogView3PlusPipelineOutput, but pipeline_output.py defines CogView3PipelineOutput. Importing the advertised lazy symbol fails.

Impact:
Public lazy imports are broken for the output type, and docs/export tooling can drift from the actual class.

Reproduction:

try:
fromdiffusers.pipelines.cogview3importCogView3PlusPipelineOutputexceptExceptionase:
print(type(e).__name__, e)
fromdiffusers.pipelines.cogview3.pipeline_outputimportCogView3PipelineOutputprint(CogView3PipelineOutput.__name__)

Relevant precedent:
Docs already autodoc the real class name at:

## CogView3PipelineOutput
[[autodoc]] pipelines.cogview3.pipeline_output.CogView3PipelineOutput

Suggested fix:

_import_structure= {"pipeline_output": ["CogView3PipelineOutput"]}
# in TYPE_CHECKING branchfrom .pipeline_outputimportCogView3PipelineOutput

Issue 2: Precomputed prompt embeddings are not expanded or dtype/device-normalized

Affected code:

ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(
prompt=prompt,
num_images_per_prompt=num_images_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
ifdo_classifier_free_guidanceandnegative_promptisNone:
negative_prompt_embeds=prompt_embeds.new_zeros(prompt_embeds.shape)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
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`."
)
negative_prompt_embeds=self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_images_per_prompt=num_images_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)

# 5. Prepare latents.
latent_channels=self.transformer.config.in_channels
latents=self.prepare_latents(
batch_size*num_images_per_prompt,
latent_channels,
height,
width,
prompt_embeds.dtype,
device,
generator,
latents,
)
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs=self.prepare_extra_step_kwargs(generator, eta)
# 7. Prepare additional timestep conditions
original_size=torch.tensor([original_size], dtype=prompt_embeds.dtype)
target_size=torch.tensor([target_size], dtype=prompt_embeds.dtype)
crops_coords_top_left=torch.tensor([crops_coords_top_left], dtype=prompt_embeds.dtype)
ifself.do_classifier_free_guidance:
original_size=torch.cat([original_size, original_size])
target_size=torch.cat([target_size, target_size])
crops_coords_top_left=torch.cat([crops_coords_top_left, crops_coords_top_left])
original_size=original_size.to(device).repeat(batch_size*num_images_per_prompt, 1)
target_size=target_size.to(device).repeat(batch_size*num_images_per_prompt, 1)
crops_coords_top_left=crops_coords_top_left.to(device).repeat(batch_size*num_images_per_prompt, 1)
# 8. Denoising loop
num_warmup_steps=max(len(timesteps) -num_inference_steps*self.scheduler.order, 0)
withself.progress_bar(total=num_inference_steps) asprogress_bar:
# for DPM-solver++
old_pred_original_sample=None
fori, tinenumerate(timesteps):
ifself.interrupt:
continue
latent_model_input=torch.cat([latents] *2) ifself.do_classifier_free_guidanceelselatents
latent_model_input=self.scheduler.scale_model_input(latent_model_input, t)
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep=t.expand(latent_model_input.shape[0])
# predict noise model_output
noise_pred=self.transformer(
hidden_states=latent_model_input,
encoder_hidden_states=prompt_embeds,
timestep=timestep,
original_size=original_size,
target_size=target_size,
crop_coords=crops_coords_top_left,
return_dict=False,
)[0]

Problem:
When prompt_embeds or negative_prompt_embeds are passed directly, CogView3 does not repeat them for num_images_per_prompt and does not cast/move them to the execution dtype/device. The latent batch is expanded, but the prompt batch is not.

Impact:
Valid diffusers usage with precomputed embeddings fails with batch mismatch, device mismatch, or dtype mismatch.

Reproduction:

importtorchfromdiffusersimportCogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModeldefmake_pipe(dtype=torch.float32):
transformer=CogView3PlusTransformer2DModel(
patch_size=2, in_channels=4, num_layers=1, attention_head_dim=4,
num_attention_heads=2, out_channels=4, text_embed_dim=8,
time_embed_dim=8, condition_dim=2, pos_embed_max_size=8, sample_size=2,
).to(dtype=dtype)
pipe=CogView3PlusPipeline(None, None, None, transformer, CogVideoXDDIMScheduler())
pipe.set_progress_bar_config(disable=True)
returnpipetry:
make_pipe()(prompt_embeds=torch.randn(1, 8, 8), num_images_per_prompt=2,
num_inference_steps=1, guidance_scale=1.0, height=16, width=16, output_type="latent")
exceptExceptionase:
print(type(e).__name__, e)
try:
make_pipe(torch.float16)(prompt_embeds=torch.randn(1, 8, 8, dtype=torch.float32),
num_inference_steps=1, guidance_scale=1.0,
height=16, width=16, output_type="latent")
exceptExceptionase:
print(type(e).__name__, e)

Relevant precedent:
CogView4 expands precomputed embeds:

ifprompt_embedsisNone:
prompt_embeds=self._get_glm_embeds(prompt, max_sequence_length, device, dtype)
seq_len=prompt_embeds.size(1)
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)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=negative_promptor""
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
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`."
)
negative_prompt_embeds=self._get_glm_embeds(negative_prompt, max_sequence_length, device, dtype)
seq_len=negative_prompt_embeds.size(1)
negative_prompt_embeds=negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
negative_prompt_embeds=negative_prompt_embeds.view(batch_size*num_images_per_prompt, seq_len, -1)

Suggested fix:

ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(...)
else:
dtype=dtypeorself.transformer.dtypeprompt_embeds=prompt_embeds.to(device=device, dtype=dtype)
bs_embed, seq_len, _=prompt_embeds.shapeprompt_embeds=prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds=prompt_embeds.view(bs_embed*num_images_per_prompt, seq_len, -1)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNoneandnegative_promptisNone:
negative_prompt_embeds=prompt_embeds.new_zeros(prompt_embeds.shape)
elifdo_classifier_free_guidanceandnegative_prompt_embedsisnotNone:
negative_prompt_embeds=negative_prompt_embeds.to(device=device, dtype=prompt_embeds.dtype)
bs_embed, seq_len, _=negative_prompt_embeds.shapenegative_prompt_embeds=negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
negative_prompt_embeds=negative_prompt_embeds.view(bs_embed*num_images_per_prompt, seq_len, -1)

Issue 3: Height/width validation ignores transformer patch size

Affected code:

# Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
defcheck_inputs(
self,
prompt,
height,
width,
negative_prompt,
callback_on_step_end_tensor_inputs,
prompt_embeds=None,
negative_prompt_embeds=None,
):
ifheight%8!=0orwidth%8!=0:
raiseValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

Problem:
The pipeline only requires dimensions divisible by 8, but CogView3 patchifies VAE latents with patch_size=2. A 24x24 image passes check_inputs, then fails inside the transformer because latent size 3x3 is not divisible by patch size.

Impact:
Users get a late internal transformer error instead of a clear pipeline validation error.

Reproduction:

importtorchfromdiffusersimportCogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModeltransformer=CogView3PlusTransformer2DModel(
patch_size=2, in_channels=4, num_layers=1, attention_head_dim=4,
num_attention_heads=2, out_channels=4, text_embed_dim=8,
time_embed_dim=8, condition_dim=2, pos_embed_max_size=8, sample_size=2,
)
pipe=CogView3PlusPipeline(None, None, None, transformer, CogVideoXDDIMScheduler())
pipe.set_progress_bar_config(disable=True)
try:
pipe(prompt_embeds=torch.randn(1, 8, 8), num_inference_steps=1,
guidance_scale=1.0, height=24, width=24, output_type="latent")
exceptExceptionase:
print(type(e).__name__, e)

Relevant precedent:
CogView4 validates 16, matching vae_scale_factor * patch_size:

ifheight%16!=0orwidth%16!=0:
raiseValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")

Suggested fix:

divisibility=self.vae_scale_factor*self.transformer.config.patch_sizeifheight%divisibility!=0orwidth%divisibility!=0:
raiseValueError(
f"`height` and `width` have to be divisible by {divisibility} but are {height} and {width}."
)

Issue 4: Attention backend selection is effectively ignored

Affected code:

from ..attentionimportAttentionMixin, FeedForward
from ..attention_processorimportAttention, CogVideoXAttnProcessor2_0

self.attn1=Attention(
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
out_dim=dim,
bias=True,
qk_norm="layer_norm",
elementwise_affine=False,
eps=1e-6,
processor=CogVideoXAttnProcessor2_0(),
)

classCogVideoXAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
query and key vectors, but does not include spatial normalization.
"""
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def__call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor|None=None,
image_rotary_emb: torch.Tensor|None=None,
) ->torch.Tensor:
text_seq_length=encoder_hidden_states.size(1)
hidden_states=torch.cat([encoder_hidden_states, hidden_states], dim=1)
batch_size, sequence_length, _=hidden_states.shape
ifattention_maskisnotNone:
attention_mask=attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask=attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
query=attn.to_q(hidden_states)
key=attn.to_k(hidden_states)
value=attn.to_v(hidden_states)
inner_dim=key.shape[-1]
head_dim=inner_dim//attn.heads
query=query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key=key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value=value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
ifattn.norm_qisnotNone:
query=attn.norm_q(query)
ifattn.norm_kisnotNone:
key=attn.norm_k(key)
# Apply RoPE if needed
ifimage_rotary_embisnotNone:
from .embeddingsimportapply_rotary_emb
query[:, :, text_seq_length:] =apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
ifnotattn.is_cross_attention:
key[:, :, text_seq_length:] =apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
hidden_states=F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
)

Problem:
CogView3Plus reuses CogVideoXAttnProcessor2_0, which calls F.scaled_dot_product_attention directly and has no _attention_backend. model.set_attention_backend(...) skips it, so backend dispatch support is a no-op for this transformer.

Impact:
Users cannot reliably select supported attention backends for CogView3Plus, and the implementation violates the review rule requiring model-local processors to route through dispatch_attention_fn.

Reproduction:

fromdiffusersimportCogView3PlusTransformer2DModelmodel=CogView3PlusTransformer2DModel(
patch_size=2, in_channels=4, num_layers=1, attention_head_dim=4,
num_attention_heads=2, out_channels=4, text_embed_dim=8,
time_embed_dim=8, condition_dim=2, pos_embed_max_size=8, sample_size=2,
)
name, processor=next(iter(model.attn_processors.items()))
print(name, processor.__class__.__name__, hasattr(processor, "_attention_backend"))
model.set_attention_backend("native")
print(hasattr(processor, "_attention_backend"), getattr(processor, "_attention_backend", None))

Relevant precedent:
QwenImage routes attention through dispatch_attention_fn:

joint_hidden_states=dispatch_attention_fn(
joint_query,
joint_key,
joint_value,
attn_mask=attention_mask,
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,

Suggested fix:
Implement a CogView3Plus-local attention processor, with _attention_backend and _parallel_config, and call dispatch_attention_fn(...) instead of F.scaled_dot_product_attention(...). If doing the full rule-compliant refactor, also define a local attention module inheriting AttentionModuleMixin.

Issue 5: Released transformer config contains an ignored key on load

Affected code:

patch_size: int=2,
in_channels: int=16,
num_layers: int=30,
attention_head_dim: int=40,
num_attention_heads: int=64,
out_channels: int=16,
text_embed_dim: int=4096,
time_embed_dim: int=512,
condition_dim: int=256,
pos_embed_max_size: int=128,
sample_size: int=128,
):
super().__init__()
self.out_channels=out_channels
self.inner_dim=num_attention_heads*attention_head_dim
# CogView3 uses 3 additional SDXL-like conditions - original_size, target_size, crop_coords
# Each of these are sincos embeddings of shape 2 * condition_dim
self.pooled_projection_dim=3*2*condition_dim
self.patch_embed=CogView3PlusPatchEmbed(
in_channels=in_channels,
hidden_size=self.inner_dim,
patch_size=patch_size,
text_hidden_size=text_embed_dim,
pos_embed_max_size=pos_embed_max_size,
)
self.time_condition_embed=CogView3CombinedTimestepSizeEmbeddings(
embedding_dim=time_embed_dim,
condition_dim=condition_dim,
pooled_projection_dim=self.pooled_projection_dim,

Problem:
The published THUDM/CogView3-Plus-3B transformer config includes pooled_projection_dim, but the model constructor does not accept it. Loading emits an unexpected-config warning even though the value matches the derived value.

Impact:
Users see a spurious load warning for the official checkpoint, and config compatibility remains noisy.

Reproduction:

fromdiffusersimportCogView3PlusTransformer2DModelconfig= {
"patch_size": 2, "in_channels": 4, "num_layers": 1,
"attention_head_dim": 4, "num_attention_heads": 2,
"out_channels": 4, "text_embed_dim": 8, "time_embed_dim": 8,
"condition_dim": 2, "pos_embed_max_size": 8, "sample_size": 2,
"pooled_projection_dim": 12,
}
model=CogView3PlusTransformer2DModel.from_config(config)
print(model.pooled_projection_dim)

Relevant precedent:
Other transformer configs expose pooled_projection_dim directly where it is serialized:

pooled_projection_dim: int=1024,
text_len: int=77,
text_len_t5: int=256,
use_style_cond_and_image_meta_size: bool=True,
):
super().__init__()
self.out_channels=in_channels*2iflearn_sigmaelsein_channels
self.num_heads=num_attention_heads
self.inner_dim=num_attention_heads*attention_head_dim
self.text_embedder=PixArtAlphaTextProjection(
in_features=cross_attention_dim_t5,
hidden_size=cross_attention_dim_t5*4,
out_features=cross_attention_dim,
act_fn="silu_fp32",
)
self.text_embedding_padding=nn.Parameter(torch.randn(text_len+text_len_t5, cross_attention_dim))
self.pos_embed=PatchEmbed(
height=sample_size,
width=sample_size,
in_channels=in_channels,
embed_dim=hidden_size,
patch_size=patch_size,
pos_embed_type=None,
)
self.time_extra_emb=HunyuanCombinedTimestepTextSizeStyleEmbedding(
hidden_size,
pooled_projection_dim=pooled_projection_dim,

Suggested fix:

def__init__(..., pooled_projection_dim: int|None=None, ...):
...
self.pooled_projection_dim=pooled_projection_dimor3*2*condition_dimself.register_to_config(pooled_projection_dim=self.pooled_projection_dim)

Issue 6: CogView3 tests/docs do not provide meaningful coverage

Affected code:

self.assertEqual(generated_image.shape, (3, 16, 16))
expected_image=torch.randn(3, 16, 16)
max_diff=np.abs(generated_image-expected_image).max()
self.assertLessEqual(max_diff, 1e10)

deftest_cogview3plus(self):
generator=torch.Generator("cpu").manual_seed(0)
pipe=CogView3PlusPipeline.from_pretrained("THUDM/CogView3Plus-3b", torch_dtype=torch.float16)
pipe.enable_model_cpu_offload(device=torch_device)
prompt=self.prompt
images=pipe(
prompt=prompt,
height=1024,
width=1024,
generator=generator,
num_inference_steps=2,
output_type="np",
)[0]
image=images[0]
expected_image=torch.randn(1, 1024, 1024, 3).numpy()
max_diff=numpy_cosine_similarity_distance(image, expected_image)
assertmax_diff<1e-3, f"Max diff is too high. got {image}"

from diffusers import CogView3PlusTransformer2DModel
transformer = CogView3PlusTransformer2DModel.from_pretrained("THUDM/CogView3Plus-3b", subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda")

Problem:
Fast tests compare against torch.randn(...) with 1e10 tolerance, so they only check shape. The slow test uses non-existent model ID THUDM/CogView3Plus-3b, loads unsupported torch.float16, and compares images[0] to a random expected array with an incompatible leading batch dimension. The model docs also use the stale checkpoint ID.

Impact:
Fast and slow tests exist, but they do not catch output regressions. The slow test cannot validate the official checkpoint as written.

Reproduction:

fromhuggingface_hubimportmodel_infoimportnumpyasnpfromnumpy.linalgimportnormformodel_idin ["THUDM/CogView3Plus-3b", "THUDM/CogView3-Plus-3B"]:
try:
print(model_id, model_info(model_id).id)
exceptExceptionase:
print(model_id, type(e).__name__)
defnumpy_cosine_similarity_distance(a, b):
similarity=np.dot(a, b) / (norm(a) *norm(b))
return1.0-similarity.mean()
try:
image=np.zeros((4, 4, 3), dtype=np.float32)
expected_image=np.random.randn(1, 4, 4, 3).astype(np.float32)
print(numpy_cosine_similarity_distance(image, expected_image))
exceptExceptionase:
print(type(e).__name__, e)

Relevant precedent:
PR fixing the pipeline example checkpoint ID:
#10211

Existing fp16 black-image duplicate:
#10343

Suggested fix:
Update docs/tests to THUDM/CogView3-Plus-3B, use torch.bfloat16 for slow inference, and replace random expectations with deterministic expected slices or stored expected statistics generated from the official checkpoint. Fast test_inference should assert a small image slice with a real tolerance, not 1e10.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions