Add chunked saving for low memory, add pre-quantized models to README

README.md

@@ -1151,9 +1151,22 @@ mflux-generate-qwen-edit \
 
 **Qwen Image Layered** is a specialized image decomposition model that separates an input image into semantically disentangled RGBA layers. This enables powerful layer-based editing workflows where each layer can be independently manipulated and recomposed—similar to working with Photoshop layers but generated automatically from any image.
 
-The model uses a custom RGBA-VAE and a Layer3D RoPE transformer architecture to understand the semantic structure of images and separate them into distinct compositional layers (foreground, background, objects, etc.).
+**Example: Using Pre-Quantized Models**
 
-**Example: Basic Image Decomposition**
+Pre-quantized models are available on HuggingFace and download automatically:
+
+```sh
+# Using 6-bit quantized model 
+mflux-generate-qwen-layered \
+  --image "input.png" \
+  --layers 4 \
+  --steps 50 \
+  --model-path zimengxiong/Qwen-Image-Layered-6bit
+```
+
+**Example: Quantizing from Full Model**
+
+If you prefer to quantize on-the-fly from the full model:
 
 ```sh
 mflux-generate-qwen-layered \
@@ -1161,41 +1174,10 @@ mflux-generate-qwen-layered \
   --layers 4 \
   --steps 50 \
   --resolution 640 \
-  --guidance 4.0 \
-  --output-dir "./layers" \
   -q 6
 ```
 
-This will generate 4 RGBA layer files (`layer_0.png`, `layer_1.png`, etc.) in the output directory.
+This downloads the full BF16 model from `Qwen/Qwen-Image-Layered` and quantizes at runtime.
-
-**Example: Fast Preview with Fewer Steps**
-
-```sh
-mflux-generate-qwen-layered \
-  --image "photo.jpg" \
-  --layers 2 \
-  --steps 10 \
-  --resolution 640 \
-  -q 6 \
-  --output-dir "./preview"
-```
-
-**Use Cases:**
-- **Layer-based editing**: Edit individual layers independently and recompose
-- **Background removal**: Extract foreground objects with transparency
-- **Image compositing**: Combine layers from multiple decomposed images
-- **Animation**: Animate individual layers for parallax or motion effects
-- **Asset extraction**: Extract clean assets from complex scenes
-
-**Tips for Qwen Image Layered:**
-1. **Resolution**: Use 640 for faster processing, 1024 for higher quality results
-2. **Number of layers**: Start with 2-4 layers; more layers require more VRAM and time
-3. **Quantization**: 6-bit quantization (`-q 6`) significantly reduces memory usage (~29GB vs ~55GB BF16)
-4. **Steps**: 50 steps provides best quality; 10-20 steps work for quick previews
-5. **Output format**: All layers are saved as RGBA PNG files with transparency
-
-⚠️ *Note: The Qwen Image Layered model requires local weights from `Qwen/Qwen-Image-Layered` (~55GB for the full model in BF16, ~29GB with 6-bit quantization). This is a research model optimized for 48GB+ Apple Silicon Macs.*
-
 ---
 
 ### 🌀 FIBO

src/mflux/models/qwen_layered/cli/save_low_memory.py

@@ -0,0 +1,101 @@
+"""
+Memory-efficient save script for Qwen-Image-Layered model.
+Loads and saves each component individually to avoid OOM.
+"""
+
+import argparse
+import gc
+import shutil
+from pathlib import Path
+
+import mlx.core as mx
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Save quantized Qwen-Image-Layered model (low memory)")
+    parser.add_argument("--path", required=True, help="Output directory for saved model")
+    parser.add_argument("-q", "--quantize", type=int, choices=[4, 6, 8], required=True, help="Quantization bits")
+    parser.add_argument("--model-path", required=True, help="Path to local model weights")
+    args = parser.parse_args()
+
+    source_path = Path(args.model_path)
+    output_path = Path(args.path)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    print("Copying tokenizer...")
+    tokenizer_src = source_path / "tokenizer"
+    tokenizer_dst = output_path / "tokenizer"
+    if tokenizer_src.exists():
+        if tokenizer_dst.exists():
+            shutil.rmtree(tokenizer_dst)
+        shutil.copytree(tokenizer_src, tokenizer_dst)
+        print("  Tokenizer copied")
+
+    from mflux.models.common.weights.loading.weight_applier import WeightApplier
+    from mflux.models.common.weights.loading.weight_loader import WeightLoader
+    from mflux.models.common.weights.saving.model_saver import ModelSaver
+    from mflux.models.qwen_layered.weights.qwen_layered_weight_definition import QwenLayeredWeightDefinition
+
+    component_info = [
+        (
+            "vae",
+            "mflux.models.qwen_layered.model.qwen_layered_vae.qwen_layered_vae",
+            "QwenLayeredVAE",
+            {"input_channels": 4, "output_channels": 4},
+            True,
+        ),
+        (
+            "text_encoder",
+            "mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder",
+            "QwenTextEncoder",
+            {},
+            False,
+        ),  # Don't quantize
+        ("transformer", "mflux.models.qwen.model.qwen_transformer.qwen_transformer", "QwenTransformer", {}, True),
+    ]
+
+    for name, module_path, class_name, kwargs, should_quantize in component_info:
+        print(f"\nProcessing {name}...")
+        gc.collect()
+        mx.metal.clear_cache() if hasattr(mx, "metal") else mx.clear_cache()
+
+        import importlib
+
+        module = importlib.import_module(module_path)
+        ComponentClass = getattr(module, class_name)
+
+        # Create component
+        component = ComponentClass(**kwargs)
+
+        # Load weights for this component only
+        print("  Loading weights...")
+        weights = WeightLoader.load(
+            weight_definition=QwenLayeredWeightDefinition,
+            model_path=str(source_path),
+        )
+
+        quant = args.quantize if should_quantize else None
+        WeightApplier.apply_and_quantize(
+            weights=weights,
+            quantize_arg=quant,
+            weight_definition=QwenLayeredWeightDefinition,
+            models={name: component},
+        )
+        mx.eval(component.parameters())
+        print(f"  Saving to {output_path / name}...")
+        ModelSaver._save_weights(str(output_path), args.quantize, component, name)
+
+        del component, weights
+        gc.collect()
+        mx.metal.clear_cache() if hasattr(mx, "metal") else mx.clear_cache()
+        print(f"  {name} done")
+
+    print(f"\nDone! Model saved to {output_path}")
+    import subprocess
+
+    result = subprocess.run(["du", "-sh", str(output_path)], capture_output=True, text=True)
+    print(f"Size: {result.stdout.strip()}")
+
+
+if __name__ == "__main__":
+    main()