diff --git a/README.md b/README.md index 84358cb..88bc511 100644 --- a/README.md +++ b/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. - -**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.* - +This downloads the full BF16 model from `Qwen/Qwen-Image-Layered` and quantizes at runtime. --- ### 🌀 FIBO diff --git a/src/mflux/models/qwen_layered/cli/save_low_memory.py b/src/mflux/models/qwen_layered/cli/save_low_memory.py new file mode 100644 index 0000000..eb877d8 --- /dev/null +++ b/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()