Skip to content

[Feature(s)] Sigma schedules, aspect ratios, shift override, MCF sampler for Z-Image #355

@terribilissimo

Description

@terribilissimo

My PR: #353

New flags/args

  • --shift: override automatic sigma shift (mu) for both schedulers
  • --mcf-max-change: MCF sampler that clamps per-step latent changes
  • --cosine / --karras / --exponential: alternative sigma schedules
  • --aspect: aspect ratio presets with auto-dimension computation
  • --saveinfo: descriptive filenames encoding generation parameters
  • All features apply to mflux-generate-z-image-turbo and mflux-generate-z-image
  • NOTE: when no fork-specific flags are used, behavior is identical to upstream

Summary

This PR adds several quality-of-life enhancements to the Z-Image and Z-Image Turbo
pipelines, ported from our zima.py CUDA/PyTorch scripts. All features are
opt-in — when none of the new flags are used, behavior is identical to upstream.

New CLI Flags

Noise Schedule Control

  • --shift <float> — Override the automatic sigma shift (mu) value.
    By default mu is computed from image dimensions. Higher values push the
    noise schedule towards higher noise levels.
  • --cosine — Smooth S-curve sigma schedule (more steps at high/low noise)
  • --karras — Karras schedule with rho=7 (concentrates steps on fine details)
  • --exponential — Log-spaced sigmas between sigma_max and sigma_min
    (Only one of the three can be used at a time; mutually exclusive group)

Sampling

  • --mcf-max-change <float> — MCF (Mean Change Factor) sampler that clamps
    per-step latent changes to prevent sudden jumps. Typical: 0.05–0.50, min/max 0.01/1.00.

Convenience

  • --aspect <preset> — Aspect ratio presets (1:1, 4:3, 3:4, 3:2, 2:3, 16:9,
    9:16, 18:9, 9:18, 21:9, 9:21). If combined with only --width or --height,
    the missing dimension is auto-computed and rounded to multiples of 16.
  • --saveinfo — Save images with descriptive filenames encoding
    timestamp, seed, steps, LoRA, scheduler, and sigma schedule info
    (convenient, you won't have to look at the EXIF or json for quick reproducibility).

Modified Files

  • src/mflux/cli/parser/parsers.py — New args, aspect ratio dict, sigma_schedule resolution; removed --metadata
  • src/mflux/models/common/config/config.pysigma_schedule property
  • src/mflux/models/common/schedulers/linear_scheduler.py_generate_base_sigmas()
  • src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py — same
  • src/mflux/models/z_image/variants/z_image.py — New params in generate_image()
  • src/mflux/models/z_image/cli/z_image_turbo_generate.py — Passthrough + saveinfo
  • src/mflux/models/z_image/cli/z_image_generate.py — Same
  • All CLI entry points (flux, flux2, qwen, z_image, fibo) — Removed export_json_metadata=args.metadata

Testing

All features tested on Apple Silicon (M-series) with both mflux-generate-z-image-turbo
and mflux-generate-z-image. Baseline output is pixel-identical to upstream when
no new flags are used. See TESTING-SCRIPT.txt for the full test matrix (18 tests).

Checklist

  • No new dependencies added
  • All sigma schedule math uses pure MLX (mx.array, mx.cos, mx.linspace)
  • Backward-compatible — default behavior unchanged
  • Passes existing tests
  • Both Z-Image and Z-Image Turbo entry points updated

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions