monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove legacy QwenImageEdit model, keep only QwenImageEditPlus (#276)

Browse Source 
Co-authored-by: Filip Strand <filip@Host-022.local>
This commit is contained in:
Filip Strand 2025-11-11 16:38:02 +01:00 committed by GitHub
parent f9768b2e1c
commit bb5392ef2d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 2 additions and 365 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pyproject.toml
View File

@ -84,7 +84,6 @@ mflux-generate-fill = "mflux.generate_fill:main"
mflux-generate-depth = "mflux.generate_depth:main"
mflux-generate-redux = "mflux.generate_redux:main"
mflux-generate-kontext = "mflux.generate_kontext:main"
mflux-generate-qwen-edit = "mflux.generate_qwen_edit:main"
mflux-generate-qwen-edit-plus = "mflux.generate_qwen_edit_plus:main"
mflux-concept = "mflux.concept:main"
mflux-concept-from-image = "mflux.concept_from_image:main"

17
src/mflux/config/model_config.py
View File

@ -89,11 +89,6 @@ class ModelConfig:
def qwen_image() -> "ModelConfig":
return AVAILABLE_MODELS["qwen-image"]
@staticmethod
@lru_cache
def qwen_image_edit() -> "ModelConfig":
return AVAILABLE_MODELS["qwen-image-edit"]
@staticmethod
@lru_cache
def qwen_image_edit_plus() -> "ModelConfig":
@ -312,18 +307,6 @@ AVAILABLE_MODELS = {
requires_sigma_shift=None,
priority=11,
),
"qwen-image-edit": ModelConfig(
aliases=["qwen-image-edit", "qwen-edit"],
model_name="Qwen/Qwen-Image-Edit",
base_model=None,
controlnet_model=None,
custom_transformer_model=None,
num_train_steps=None,
max_sequence_length=None,
supports_guidance=None,
requires_sigma_shift=None,
priority=12,
),
"qwen-image-edit-plus": ModelConfig(
aliases=["qwen-image-edit-plus", "qwen-edit-plus", "qwen-edit-2509"],
model_name="Qwen/Qwen-Image-Edit-2509",

65
src/mflux/generate_qwen_edit.py
View File

@ -1,65 +0,0 @@
from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
def main():
# 0. Parse command line arguments
parser = CommandLineParser(description="Generate an image using Flux Kontext with image conditioning.")
parser.add_general_arguments()
parser.add_model_arguments(require_model_arg=False)
parser.add_lora_arguments()
parser.add_image_generator_arguments(supports_metadata_config=True)
parser.add_image_to_image_arguments(required=True)
parser.add_output_arguments()
args = parser.parse_args()
# 0. Set default guidance value if not provided by user
if args.guidance is None:
args.guidance = ui_defaults.GUIDANCE_SCALE_KONTEXT
# 1. Load the model
qwen = QwenImageEdit(
quantize=args.quantize,
local_path=args.path,
lora_paths=args.lora_paths,
lora_scales=args.lora_scales,
)
# 2. Register callbacks
memory_saver = CallbackManager.register_callbacks(args=args, model=qwen)
try:
for seed in args.seed:
# 3. Generate an image for each seed value
image = qwen.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,
width=args.width,
guidance=args.guidance,
image_path=args.image_path,
),
)
# 4. Save the image
output_path = Path(args.output.format(seed=seed))
image.save(path=output_path, export_json_metadata=args.metadata)
except (StopImageGenerationException, PromptFileReadError) as exc:
print(exc)
finally:
if memory_saver:
print(memory_saver.memory_stats())
if __name__ == "__main__":
main()

3
src/mflux/models/qwen/variants/edit/__init__.py
View File

@ -1,5 +1,4 @@
# Qwen Image Edit variant
from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
from mflux.models.qwen.variants.edit.qwen_image_edit_plus import QwenImageEditPlus
__all__ = ["QwenImageEdit", "QwenImageEditPlus"]
__all__ = ["QwenImageEditPlus"]

263
src/mflux/models/qwen/variants/edit/qwen_image_edit.py
View File

@ -1,263 +0,0 @@
import math
import mlx.core as mx
from mlx import nn
from PIL import Image
from tqdm import tqdm
from mflux.callbacks.callbacks import Callbacks
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.error.exceptions import StopImageGenerationException
from mflux.latent_creator.latent_creator import LatentCreator
from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE
from mflux.models.qwen.qwen_edit_initializer import QwenImageEditInitializer
from mflux.models.qwen.variants.edit.utils.qwen_edit_util import QwenEditUtil
from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.generated_image import GeneratedImage
from mflux.post_processing.image_util import ImageUtil
class QwenImageEdit(nn.Module):
vae: QwenVAE
transformer: QwenTransformer
text_encoder: QwenTextEncoder
def __init__(
self,
quantize: int | None = None,
local_path: str | None = None,
lora_paths: list[str] | None = None,
lora_scales: list[float] | None = None,
lora_names: list[str] | None = None,
lora_repo_id: str | None = None,
):
super().__init__()
QwenImageEditInitializer.init(
qwen_model=self,
model_config=ModelConfig.qwen_image_edit(),
quantize=quantize,
local_path=local_path,
lora_paths=lora_paths,
lora_scales=lora_scales,
lora_names=lora_names,
lora_repo_id=lora_repo_id,
)
def generate_image(
self,
seed: int,
prompt: str,
config: Config,
negative_prompt: str | None = None,
) -> GeneratedImage:
runtime_config, vl_width, vl_height = self._compute_dimensions(config)
timesteps = runtime_config.scheduler.timesteps
time_steps = tqdm(range(len(timesteps)))
# 1. Create initial latents
latents = LatentCreator.create(
seed=seed,
height=runtime_config.height,
width=runtime_config.width,
)
# 2. Encode the prompt
prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = (
self._debug_encode_prompts_with_image(
prompt=prompt,
negative_prompt=negative_prompt,
image_path=runtime_config.image_path,
runtime_config=runtime_config,
vl_width=vl_width,
vl_height=vl_height,
)
)
# 3. Generate image conditioning latents
image_path = str(runtime_config.image_path) if runtime_config.image_path else None
static_image_latents, qwen_image_ids, cond_h_patches, cond_w_patches, num_images = (
QwenEditUtil.create_image_conditioning_latents(
vae=self.vae,
height=runtime_config.height,
width=runtime_config.width,
image_paths=image_path,
vl_width=vl_width,
vl_height=vl_height,
)
)
# (Optional) Call subscribers for beginning of loop
Callbacks.before_loop(
seed=seed,
prompt=prompt,
latents=latents,
config=runtime_config,
)
for t in time_steps:
try:
# 4.t Concatenate the updated latents with the static image latents
hidden_states = mx.concatenate([latents, static_image_latents], axis=1)
hidden_states_neg = mx.concatenate([latents, static_image_latents], axis=1) # noqa: F841
# 5.t Predict the noise
noise = self.transformer(
t=t,
config=runtime_config,
hidden_states=hidden_states,
encoder_hidden_states=prompt_embeds,
encoder_hidden_states_mask=prompt_mask,
qwen_image_ids=qwen_image_ids,
cond_image_grid=(1, cond_h_patches, cond_w_patches),
)[:, : latents.shape[1]]
noise_negative = self.transformer(
t=t,
config=runtime_config,
hidden_states=hidden_states_neg,
encoder_hidden_states=negative_prompt_embeds,
encoder_hidden_states_mask=negative_prompt_mask,
qwen_image_ids=qwen_image_ids,
cond_image_grid=(1, cond_h_patches, cond_w_patches),
)[:, : latents.shape[1]]
guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance)
# 6.t Take one denoise step
latents = runtime_config.scheduler.step(
model_output=guided_noise,
timestep=t,
sample=latents,
)
# (Optional) Call subscribers in-loop
Callbacks.in_loop(
t=t,
seed=seed,
prompt=prompt,
latents=latents,
config=runtime_config,
time_steps=time_steps,
)
# (Optional) Evaluate to enable progress tracking
mx.eval(latents)
except KeyboardInterrupt: # noqa: PERF203
Callbacks.interruption(
t=t,
seed=seed,
prompt=prompt,
latents=latents,
config=runtime_config,
time_steps=time_steps,
)
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(timesteps)}")
# (Optional) Call subscribers after loop
Callbacks.after_loop(
seed=seed,
prompt=prompt,
latents=latents,
config=runtime_config,
)
# 7. Decode the latent array and return the image
latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
config=runtime_config,
seed=seed,
prompt=prompt,
quantization=self.bits,
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
image_path=runtime_config.image_path,
generation_time=time_steps.format_dict["elapsed"],
negative_prompt=negative_prompt,
)
def _debug_encode_prompts_with_image(
self,
prompt: str,
negative_prompt: str,
image_path: str,
runtime_config,
vl_width: int | None = None,
vl_height: int | None = None,
) -> tuple[mx.array, mx.array, mx.array, mx.array]:
tokenizer = self.qwen_vl_tokenizer
image = Image.open(image_path).convert("RGB")
pos_input_ids, pos_attention_mask, pos_pixel_values, pos_image_grid_thw = tokenizer.tokenize_with_image(
prompt, image, vl_width=vl_width, vl_height=vl_height
)
pos_hidden_states = self.qwen_vl_encoder(
input_ids=pos_input_ids,
attention_mask=pos_attention_mask,
pixel_values=pos_pixel_values,
image_grid_thw=pos_image_grid_thw,
)
neg_prompt = negative_prompt if negative_prompt is not None else ""
neg_input_ids, neg_attention_mask, neg_pixel_values, neg_image_grid_thw = tokenizer.tokenize_with_image(
neg_prompt, image, vl_width=vl_width, vl_height=vl_height
)
neg_hidden_states = self.qwen_vl_encoder(
input_ids=neg_input_ids,
attention_mask=neg_attention_mask,
pixel_values=neg_pixel_values,
image_grid_thw=neg_image_grid_thw,
)
return (
pos_hidden_states[0].astype(mx.float16), # prompt_embeds
pos_hidden_states[1].astype(mx.float16), # prompt_mask
neg_hidden_states[0].astype(mx.float16), # negative_prompt_embeds
neg_hidden_states[1].astype(mx.float16), # negative_prompt_mask
)
def _compute_dimensions(
self,
config: Config,
) -> tuple[RuntimeConfig, int, int]:
image = ImageUtil.load_image(config.image_path).convert("RGB")
image_size = image.size
target_area = 1024 * 1024
ratio = image_size[0] / image_size[1]
calculated_width = math.sqrt(target_area * ratio)
calculated_height = calculated_width / ratio
calculated_width = round(calculated_width / 32) * 32
calculated_height = round(calculated_height / 32) * 32
use_height = config.height or calculated_height
use_width = config.width or calculated_width
vae_scale_factor = 8
multiple_of = vae_scale_factor * 2
use_width = use_width // multiple_of * multiple_of
use_height = use_height // multiple_of * multiple_of
final_config = Config(
height=use_height,
width=use_width,
image_path=config.image_path,
num_inference_steps=config.num_inference_steps,
guidance=config.guidance,
scheduler=config.scheduler_str,
)
runtime_config = RuntimeConfig(final_config, self.model_config)
vl_width = calculated_width
vl_height = calculated_height
return runtime_config, int(vl_width), int(vl_height)

18
tests/image_generation/test_generate_image_edit.py
View File

@ -1,6 +1,6 @@
from mflux.config.model_config import ModelConfig
from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
from mflux.models.qwen.variants.edit import QwenImageEdit, QwenImageEditPlus
from mflux.models.qwen.variants.edit import QwenImageEditPlus
from tests.image_generation.helpers.image_generation_edit_test_helper import ImageGeneratorEditTestHelper
@ -20,22 +20,6 @@ class TestImageGeneratorEdit:
image_path="reference_upscaled.png",
)
def test_image_generation_qwen_edit(self):
ImageGeneratorEditTestHelper.assert_matches_reference_image(
reference_image_path="reference_qwen_edit.png",
output_image_path="output_qwen_edit.png",
model_class=QwenImageEdit,
model_config=ModelConfig.qwen_image_edit(),
steps=20,
seed=4869845,
height=384,
width=640,
guidance=2.5,
quantize=6,
prompt="Make the hand fistbump the camera instead of showing a flat palm",
image_path="reference_upscaled.png",
)
def test_image_generation_qwen_edit_plus(self):
ImageGeneratorEditTestHelper.assert_matches_reference_image(
reference_image_path="reference_qwen_edit_plus.png",

BIN
tests/resources/reference_qwen_edit.png
View File

Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 216 KiB