Use float32 Euler updates for exact five-step sampling

README.md

@@ -19,7 +19,7 @@ tags:
 
 # Qwen-Image-2.1-P-Turbo
 
-**Improved using Qwen.** A five-step LoRA adapter for [Qwen Image 2.1](https://huggingface.co/Qwen/Qwen-Image-2.1). The adapter weights are a modification of the base model; its license applies. The `_comfyui_bf16.safetensors` file is for ComfyUI. Use LoRA strength 1.0 and sigma schedule `1 → 13/14 → 6/7 → 2/3 → 0.4 → 0`.
+**Improved using Qwen.** A five-step LoRA adapter for [Qwen Image 2.1](https://huggingface.co/Qwen/Qwen-Image-2.1). The adapter weights are a modification of the base model; its license applies. The `_comfyui_bf16.safetensors` file is for ComfyUI. Use LoRA strength 1.0, float32 Euler updates, and sigma schedule `1 → 13/14 → 6/7 → 2/3 → 0.4 → 0`.
 
 ## Quick Start
 
@@ -33,6 +33,16 @@ hf auth login  # required while this repository is private
 import torch
 from PIL import Image
 from diffusers import FlowMatchEulerDiscreteScheduler, QwenImage21Pipeline
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteSchedulerOutput
+
+class Float32EulerScheduler(FlowMatchEulerDiscreteScheduler):
+    def step(self, model_output, timestep, sample, return_dict=True, **kwargs):
+        if self.step_index is None:
+            self._init_step_index(timestep)
+        dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index]
+        prev = (sample.float() + dt * model_output.float()).to(model_output.dtype)
+        self._step_index += 1
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev) if return_dict else (prev,)
 
 pipe = QwenImage21Pipeline.from_pretrained(
     "Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16
@@ -41,7 +51,7 @@ pipe.load_lora_weights(
     "PrunaAI/Qwen-Image-2.1-P-Turbo",
     weight_name="p_qwen_image_2.1_5step_v1.0.safetensors",
 )
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+pipe.scheduler = Float32EulerScheduler.from_config(
     pipe.scheduler.config,
     use_dynamic_shifting=False,
     shift=1.0,
@@ -53,7 +63,7 @@ sigmas = [1.0, 13 / 14, 6 / 7, 2 / 3, 0.4]  # terminal 0 is appended
 image = pipe(
     prompt='A neon shop sign that reads "QWEN IMAGE 2.1", rainy night',
     width=1024, height=1024,
-    num_inference_steps=5, sigmas=sigmas, true_cfg_scale=1.0,
+    num_inference_steps=5, sigmas=sigmas, true_cfg_scale=1.0, use_kv_cache=False,
     generator=torch.Generator("cuda").manual_seed(42),
 ).images[0]
 image.save("t2i_example.png")
@@ -63,7 +73,7 @@ input_image = Image.open("input.png").convert("RGB")
 image = pipe(
     prompt="Change the background to a sunset beach",
     image=input_image,
-    num_inference_steps=5, sigmas=sigmas, true_cfg_scale=1.0,
+    num_inference_steps=5, sigmas=sigmas, true_cfg_scale=1.0, use_kv_cache=False,
     generator=torch.Generator("cuda").manual_seed(42),
 ).images[0]
 image.save("edit_example.png")