fix bug for janus model image generation

[kaixuanliu]

Fix issue in #44792. @zucchini-nlp @ydshieh pls help review, thx!

[zucchini-nlp]

left a couple q

[kaixuanliu]

@zucchini-nlp ,Hi, can you help review it again? Thx!!

[ydshieh]

Investigation notes [Just for the record, no need to read]

9a6df2ce is the last commit where the test passes completely (no fixes needed).

Observations from bisecting the regression, in commit order:

bdaddb6f — Generation works (after removing 2 unrelated asserts at the top of the test, or equivalently with 9daee2e8), but produced output differs from expected values → test assertion mismatch only.

a81e04a9 — After removing the 2 asserts, generation now crashes:

TypeError: '>' not supported between instances of 'int' and 'NoneType'

at max(generation_config.max_length, num_image_tokens + seq_len) — generation_config.max_length became None.

9daee2e8 / 2877e4e2 — Same max_length=None crash, no longer need to remove the 2 asserts to reproduce it.

93d7affd ("Generation config boolean defaults #43000") — New crash:

TypeError: repeat_interleave() received an invalid combination of arguments - got (NoneType, dim=int)

generation_config.num_return_sequences became None, passed as expand_size=None into _expand_inputs_for_generation.

Current main — Same expand_size=None / repeat_interleave(None) crash.

---

Experimenting with the PR fix:

• PR fix applied, but without is_first_iteration=True: the expand_size crash is gone, but generation now fails deep in RoPE:

  CUDA error: device-side assert triggered
  

  at apply_rotary_pos_emb — position_ids are computed incorrectly, causing an out-of-bounds RoPE index error.

• Full PR fix (with is_first_iteration=True): generation completes successfully. Only remaining issue is the expected output values in the test needing to be updated, which the PR handles.

So is_first_iteration=True is necessary to work around a behavioral change introduced somewhere between bdaddb6f and current main in prepare_inputs_for_generation (likely 3c52b78 — "Always pass full input_ids in prepare_inputs_for_generation"), which changed how input_ids are sliced when use_cache=True. Without is_first_iteration=True, the base class slices input_tokens (the full prompt) to 1 token, producing wrong position_ids.

[ydshieh]

ok, the failure, if we don't include the change is_first_iteration, comes from

421c7f6 [core] 🚨 Completely remove cache positions (#44181)

I ping him internally

[ydshieh]

hi, on which hardware you get this value for cuda? On A10, I get

([ 2567, 6155, 6155, 250, 15131, 15797, 15453, 12190, 3351, 10803,

[kaixuanliu]

I use A100 with torch version 2.11.0+cu128. You can adjust the expected token to adapt to your CI env.

[ydshieh]

This image test for janus is broken for so long, with different errors introduced over several commits (and some of them are resolved).

This is_first_iteration=True not only fixes the crash issue (I didn't find the root commit for it yet) but also bring the actual outputs back to match the expected outputs (which should have been updated in Default auto (#42805) ).

This fix is thus valid .

[ydshieh]

so short (or long) history

• Default auto (Default auto 🚨 🚨  #42805) --> we should have updated the expected output but we didn't (dtype change)
• Prefill-related logic in input preparation for generation (Prefill-related logic in input preparation for generation #42088) --> This changes the expected output further, which requires is_first_iteration=True to bring it back to normal
• Completely remove cache positions ([core] 🚨 Completely remove cache positions #44181) --> this will cause the crash, and is_first_iteration=True will avoid it.

[Cyrilvallez]

If we are forced to use inputs_embeds, then this fix is correct - otherwise it would indeed use input_ids without is_first_iteration. Is this expected to create inputs_embeds like that @zucchini-nlp ? Can't we let the model do it in forward from input_ids?

[ydshieh]

I will put a comment along the code, but would like to move on by merge for now.

Use inputs_embeds doesn't seem a bad thing here (no need to recompute stuff in the for loop).
I do agree that, it's strange using input_ids won't work (giving wrong value part, the crash part I have no idea).

[zucchini-nlp]

not doable for this model, because embeddings in image-generation mode are obtained via embed+pooling, later the lm head is also a bit different. In text gen generation is simple lm-style tho, which is why we have early exit a few lines above

https://github.com/kaixuanliu/transformers/blob/e634aa1bcb43e81bc12e4977bf2a673838ef7836/src/transformers/models/janus/modeling_janus.py#L1367

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: janus

[ydshieh]

run-slow: janus

[github-actions]

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/janus"]
quantizations: []

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	b6c6bbad	workflow commit (merge commit)
PR	96779dcf	branch commit (from PR)
main	9914a364	base commit (on main)

✅ No failing test specific to this PR 🎉 👏 !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.