Fix Qwen Image and Qwen Image 2.1 Generate and dtype

docs/qwen_image.md

@@ -92,15 +92,17 @@ generate(
     height=None,
     width=None,
     output_type="image",
+    **conditioning,
 )
 ```
 
 | Arg | Default | Meaning |
 |---|---|---|
 | `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 |

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)

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

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(
             {