Enable CFG for Qwen-Image

Author: ciaranbor

src/exo/worker/engines/image/models/qwen/config.py

@@ -24,7 +24,7 @@
     default_steps={"low": 10, "medium": 25, "high": 50},
     num_sync_steps_factor=0.125,  # ~3 sync steps for medium (30 steps)
     uses_attention_mask=True,  # Qwen uses encoder_hidden_states_mask
-    guidance_scale=None,  # Set to None or < 1.0 to disable CFG
+    guidance_scale=3.5,  # Set to None or < 1.0 to disable CFG
 )
 
 # Qwen-Image-Edit uses the same architecture but different processing pipeline
@@ -45,5 +45,5 @@
     default_steps={"low": 10, "medium": 25, "high": 50},
     num_sync_steps_factor=0.125,
     uses_attention_mask=True,
-    guidance_scale=None,
+    guidance_scale=3.5,
 )

src/exo/worker/engines/image/pipeline/block_wrapper.py

@@ -97,6 +97,14 @@ def reset_cache(self) -> None:
         """Reset the KV cache. Call at the start of a new generation."""
         self._kv_cache = None
 
+    def set_encoder_mask(self, mask: mx.array | None) -> None:  # noqa: B027
+        """Set the encoder hidden states mask for attention.
+
+        Override in subclasses that use attention masks (e.g., Qwen).
+        Default is a no-op for models that don't use masks (e.g., Flux).
+        """
+        del mask  # Unused in base class
+
     def __call__(
         self,
         hidden_states: mx.array,

src/exo/worker/engines/image/pipeline/runner.py

@@ -420,6 +420,13 @@ def _forward_pass(
         # Ensure wrappers are initialized (lazy - needs text_seq_len)
         self._ensure_wrappers(text_seq_len, encoder_hidden_states_mask)
 
+        # Update masks on all joint block wrappers for this pass.
+        # This is necessary for CFG where we run positive and negative passes
+        # with different masks. Qwen uses masks; Flux doesn't.
+        if self.joint_block_wrappers and encoder_hidden_states_mask is not None:
+            for wrapper in self.joint_block_wrappers:
+                wrapper.set_encoder_mask(encoder_hidden_states_mask)
+
         scaled_latents = config.scheduler.scale_model_input(latents, t)
 
         # For edit mode: concatenate with conditioning latents