diff --git a/docs/qwen_image.md b/docs/qwen_image.md index 9e9722d1..dc76ed7a 100644 --- a/docs/qwen_image.md +++ b/docs/qwen_image.md @@ -92,6 +92,7 @@ generate( height=None, width=None, output_type="image", + **conditioning, ) ``` @@ -99,8 +100,9 @@ generate( |---|---|---| | `input_ids` | required | ChatML-templated token ids, `**tokenizer(prompts)` | | `attention_mask` | `None` | padding mask from the tokenizer | -| `negative_input_ids` | `None` | tokenized negative prompt for true CFG; a space / empty prompt when unset | +| `negative_input_ids` | `None` | tokenized negative prompt for true CFG; a templated space prompt when unset | | `negative_attention_mask` | `None` | mask for the negative ids | +| `negative_prompt` | `None` | optional text (or one string per image) passed via `**conditioning`; tokenized when negative ids are unset | | `num_inference_steps` | `None` | scheduler steps; the repo's `generate_args` (50) when unset | | `guidance_scale` | `None` | true CFG strength; `generate_args` (4.0) when unset, `<= 1` disables it | | `seed` | `None` | seed for the initial latent, reproducible per backend | diff --git a/zeromodels/models/qwen_image/qwen_image_model.py b/zeromodels/models/qwen_image/qwen_image_model.py index 4ff2d36d..ff7bda67 100644 --- a/zeromodels/models/qwen_image/qwen_image_model.py +++ b/zeromodels/models/qwen_image/qwen_image_model.py @@ -5,6 +5,7 @@ from keras import layers, ops from zeromodels.base import BaseDiffusion, BaseModel +from zeromodels.base.base_attention import with_model_attn_implementation from zeromodels.base.base_mixin import inference_scope from zeromodels.base.base_scheduler import ( FlowMatchEulerDiscreteScheduler, @@ -620,12 +621,29 @@ def _latent_side(self, height=None, width=None): w = 2 * (int(width) // scale) return h, w + def _negative_prompt_inputs(self, batch, prompt=None): + """Tokenize the templated negative prompt used for true CFG.""" + from .qwen_image_tokenizer import QwenImageTokenizer + + if prompt is None: + prompts = [" "] * batch + elif isinstance(prompt, str): + prompts = [prompt] * batch + else: + prompts = list(prompt) + if len(prompts) != batch: + raise ValueError("negative_prompt must have one entry per image") + tokenizer = getattr(self, "_negative_prompt_tokenizer", None) + if tokenizer is None: + tokenizer = QwenImageTokenizer.from_weights("zeromodels/qwen-image") + self._negative_prompt_tokenizer = tokenizer + return tokenizer(prompts) + def unconditional_ids(self, batch): - # Empty / space negative prompt is encoded by the tokenizer template; - # here build a minimal pad row the task replaces via encode_negative_prompt. - length = self.config.text_config.max_seq_len - row = [self.config.pad_token_id] * length - return ops.convert_to_tensor([row] * batch, dtype="int32") + """Token IDs for a templated space prompt, matching Diffusers CFG.""" + return ops.convert_to_tensor( + self._negative_prompt_inputs(batch)["input_ids"], dtype="int32" + ) def encode_prompt(self, input_ids, attention_mask=None, **conditioning): """Encode ChatML-templated token ids → truncated prompt embeds + mask. @@ -668,6 +686,7 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): "encoder_hidden_states_mask": ops.convert_to_tensor(out_mask), } + @with_model_attn_implementation def predict_noise(self, latents, timesteps, embeddings): return self.transformer( { @@ -757,24 +776,17 @@ def generate( with inference_scope(): embeddings = self.encode_prompt(input_ids, attention_mask, **conditioning) - do_cfg = guidance_scale > 1.0 and ( - negative_input_ids is not None - or conditioning.get("negative_prompt") is not None - ) - # Diffusers enables true CFG when a negative prompt is provided. - if guidance_scale > 1.0 and negative_input_ids is None: - # Encode empty/space negative via unconditional_ids path. - neg_ids = self.unconditional_ids(batch) - neg_mask = ops.ones_like(neg_ids) - # Prefer caller-supplied negative when present. - uncond = self.encode_prompt(neg_ids, neg_mask) - do_cfg = True - elif negative_input_ids is not None: + do_cfg = guidance_scale > 1.0 + if do_cfg: + if negative_input_ids is None: + negative_inputs = self._negative_prompt_inputs( + batch, conditioning.get("negative_prompt") + ) + negative_input_ids = negative_inputs["input_ids"] + negative_attention_mask = negative_inputs["attention_mask"] uncond = self.encode_prompt(negative_input_ids, negative_attention_mask) - do_cfg = guidance_scale > 1.0 else: uncond = None - do_cfg = False # Diffusers calculate_shift: resolution-dependent flow-match mu h, w = self._latent_side(height, width) diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_layers.py b/zeromodels/models/qwen_image_21/qwen_image_21_layers.py index 25cbad0c..d1c1e290 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_layers.py +++ b/zeromodels/models/qwen_image_21/qwen_image_21_layers.py @@ -393,6 +393,7 @@ class QwenImage21Attention(layers.Layer): def __init__(self, dim, heads, dim_head, eps=NORM_EPS, module_path=None, **kwargs): if module_path is not None: kwargs.setdefault("name", safe_name(module_path)) + kwargs.setdefault("autocast", False) super().__init__(**kwargs) self.dim = dim self.heads = heads @@ -483,6 +484,8 @@ def __init__( ): if module_path is not None: kwargs.setdefault("name", safe_name(module_path)) + # Keep the float32 rotary argument intact on its way to attention. + kwargs.setdefault("autocast", False) super().__init__(**kwargs) self.dim = dim self.num_attention_heads = num_attention_heads diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_model.py b/zeromodels/models/qwen_image_21/qwen_image_21_model.py index d9e97445..b8643806 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_model.py +++ b/zeromodels/models/qwen_image_21/qwen_image_21_model.py @@ -5,6 +5,7 @@ from keras import layers, ops from zeromodels.base import BaseDiffusion, BaseModel, CausalMask +from zeromodels.base.base_attention import with_model_attn_implementation from zeromodels.base.base_mixin import inference_scope from zeromodels.base.base_scheduler import ( FlowMatchEulerDiscreteScheduler, @@ -345,7 +346,34 @@ def __init__( joint = ops.concatenate([encoder, hidden], axis=1) - rotary = ops.convert_to_tensor(rope_angles) + # The graph reserves ``text_seq_len`` slots, but Diffusers places image + # tokens immediately after the actual (batch-padded) prompt. Only the + # frame axis of image RoPE depends on that starting position. + frame_dims = axes_dims_rope[0] // 2 + frame_freqs = np.power( + 10000.0, + -np.arange(0, axes_dims_rope[0], 2, dtype=np.float64) / axes_dims_rope[0], + ).astype(np.float32) + prompt_len = ops.max(ops.sum(enc_mask_in, axis=1)) + frame_angles = ops.cast(prompt_len, "float32") * ops.convert_to_tensor( + frame_freqs, dtype="float32" + ) + image_angles = ops.concatenate( + [ + ops.broadcast_to(frame_angles, (img_seq, frame_dims)), + ops.convert_to_tensor( + rope_angles[text_seq_len:, frame_dims:], dtype="float32" + ), + ], + axis=-1, + ) + rotary = ops.concatenate( + [ + ops.convert_to_tensor(rope_angles[:text_seq_len], dtype="float32"), + image_angles, + ], + axis=0, + ) attn_mask = ops.convert_to_tensor(attn_mask_np) target_mask_t = ops.convert_to_tensor(target_token_mask) text_pos = np.flatnonzero(~image_pad_mask).astype(np.int32) @@ -730,6 +758,7 @@ def encode_prompt(self, input_ids, attention_mask=None, **conditioning): "encoder_hidden_states_mask": ops.convert_to_tensor(out_mask), } + @with_model_attn_implementation def predict_noise(self, latents, timesteps, embeddings): return self.transformer( {