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

Co-authored-by: Filip Strand <filip@Host-022.local>

pyproject.toml

@@ -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"

src/mflux/config/model_config.py

@@ -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",

src/mflux/generate_qwen_edit.py

@@ -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()

src/mflux/models/qwen/variants/edit/__init__.py

@@ -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"]

src/mflux/models/qwen/variants/edit/qwen_image_edit.py

@@ -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)

tests/image_generation/test_generate_image_edit.py

@@ -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",