For Z-Image-Turbo, make batch size > 1 work.

[hinablue]

Test on my Mac and work fine (Yeap, I also make the MPS work, but this PR only for Lokr and BS > 1).
These changes should not affect the training of other models.

[hinablue]

Quantize make the result very bad, I need more test.

[hinablue]

Close it for test.

After 1,000 steps of testing, I found the problem was with generating the sample images; the results for the sample images were terrible. The downloaded Lokr model, however, was usable. I think I need to find the cause.

[hinablue]

I updated lokr.py to make it usable for generating example images for Z-Image-Turbo Lokr. And Lokr works correctly after training.

New Logical Flow

1. Dequantization: Remains the same.
2. Base Output: Directly calls self.org_forward(x) to get the output from the original layer.
3. Delta Weight: Calculates only the LoKR delta matrix $\Delta W$.
4. Delta Output: Calculates $x \cdot \Delta W$.
5. Output Merging: Adds the two together using the formula: $Output = BaseOutput + (DeltaOutput \cdot multiplier)$.

Pros and Cons

• Pros: Excellent compatibility (works regardless of how the base model is optimized/quantized); lower memory pressure.
• Cons: Slightly slower computation speed (requires two matrix multiplications).

[hinablue]

Lokr support just update lycoris_lora==3.4.0 🙂

[hinablue]

Rollback all core changes and focus only on z_image.py to make it compatible with BS > 1.

[Yudelei]

I tested this branch on my RTX 5070 Ti (using only Z-Image Turbo + Base model), and batch size=2 worked perfectly—no errors, training ran smoothly.

Thank you very much for your contribution—this fixes a really useful feature for many of us!

[TigonTX]

How do I apply this fix to my config?

[Yudelei]

How do I apply this fix to my config?

try
git checkout pr-649

[jaretburkett]

This issue should be resolved now.