fe6936449d
This PR adds support for the Qwen-Image-Layered model, which decomposes an input image into semantically disentangled RGBA layers for layer-based editing workflows. ## Features - New CLI command: \`mflux-generate-qwen-layered\` - Decomposes images into N RGBA layers (default 4) - Supports 6-bit quantization for ~29GB memory usage (vs 55GB BF16) - Resolution buckets: 640 and 1024 ## Architecture - RGBA-VAE (4-channel) with 3D temporal convolutions for layer handling - Layer3D RoPE: 3D positional encoding [layer, height, width] - Uses base QwenTransformer with extended RoPE for multi-layer sequences - Condition image encoded as layer_index=-1 for proper decomposition ## New Files - \`src/mflux/models/qwen_layered/\` - Full model implementation - \`model/qwen_layered_vae/\` - RGBA-VAE encoder/decoder - \`model/qwen_layered_transformer/\` - Layer3D RoPE - \`weights/\` - Weight mapping and definitions - \`variants/i2l/\` - Image-to-Layers pipeline - \`cli/\` - Command-line interface ## Usage \`\`\`sh mflux-generate-qwen-layered \\ --image input.png \\ --layers 4 \\ --steps 50 \\ -q 6 \\ --output-dir ./layers \`\`\` ## Documentation Added comprehensive documentation to README.md including: - TOC entry - CLI argument reference - Usage examples and tips - Memory requirements Tested on M4 Max 48GB with 6-bit quantization.
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
import mlx.core as mx
|
|
|
|
from mflux.callbacks.callback_registry import CallbackRegistry
|
|
from mflux.models.common.config import ModelConfig
|
|
from mflux.models.common.tokenizer import TokenizerLoader
|
|
from mflux.models.common.weights.loading.loaded_weights import LoadedWeights
|
|
from mflux.models.common.weights.loading.weight_applier import WeightApplier
|
|
from mflux.models.common.weights.loading.weight_loader import WeightLoader
|
|
from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
|
|
from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer # Use base transformer
|
|
from mflux.models.qwen_layered.model.qwen_layered_vae.qwen_layered_vae import QwenLayeredVAE
|
|
from mflux.models.qwen_layered.weights.qwen_layered_weight_definition import QwenLayeredWeightDefinition
|
|
|
|
|
|
class QwenLayeredInitializer:
|
|
"""Initializer for Qwen-Image-Layered model."""
|
|
|
|
@staticmethod
|
|
def init(
|
|
model,
|
|
model_config: ModelConfig,
|
|
quantize: int | None,
|
|
model_path: str | None = None,
|
|
lora_paths: list[str] | None = None,
|
|
lora_scales: list[float] | None = None,
|
|
) -> None:
|
|
"""Initialize the Qwen-Image-Layered model."""
|
|
path = model_path if model_path else model_config.model_name
|
|
|
|
QwenLayeredInitializer._init_config(model, model_config)
|
|
weights = QwenLayeredInitializer._load_weights(path)
|
|
QwenLayeredInitializer._init_tokenizers(model, path)
|
|
QwenLayeredInitializer._init_models(model)
|
|
QwenLayeredInitializer._apply_weights(model, weights, quantize)
|
|
# Note: LoRA not supported yet for layered model
|
|
|
|
@staticmethod
|
|
def _init_config(model, model_config: ModelConfig) -> None:
|
|
model.model_config = model_config
|
|
model.lora_paths = None
|
|
model.lora_scales = None
|
|
model.prompt_cache = {}
|
|
model.callbacks = CallbackRegistry()
|
|
model.bits = None
|
|
|
|
@staticmethod
|
|
def _load_weights(path: str) -> LoadedWeights:
|
|
return WeightLoader.load(
|
|
weight_definition=QwenLayeredWeightDefinition,
|
|
model_path=path,
|
|
)
|
|
|
|
@staticmethod
|
|
def _init_tokenizers(model, path: str) -> None:
|
|
model.tokenizers = TokenizerLoader.load_all(
|
|
definitions=QwenLayeredWeightDefinition.get_tokenizers(),
|
|
model_path=path,
|
|
)
|
|
|
|
@staticmethod
|
|
def _init_models(model) -> None:
|
|
"""Initialize model components."""
|
|
model.vae = QwenLayeredVAE(input_channels=4, output_channels=4)
|
|
model.transformer = QwenTransformer() # Use base transformer
|
|
model.text_encoder = QwenTextEncoder()
|
|
|
|
@staticmethod
|
|
def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None:
|
|
"""Apply weights and optionally quantize."""
|
|
model.bits = WeightApplier.apply_and_quantize(
|
|
weights=weights,
|
|
quantize_arg=quantize,
|
|
weight_definition=QwenLayeredWeightDefinition,
|
|
models={
|
|
"vae": model.vae,
|
|
"transformer": model.transformer,
|
|
"text_encoder": model.text_encoder,
|
|
},
|
|
)
|
|
|
|
# Evaluate to load weights into memory
|
|
mx.eval(model.parameters())
|