diff --git a/pyproject.toml b/pyproject.toml index 3dfeb0d..5ccf476 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/mflux/config/model_config.py b/src/mflux/config/model_config.py index e400e25..d063055 100644 --- a/src/mflux/config/model_config.py +++ b/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", diff --git a/src/mflux/generate_qwen_edit.py b/src/mflux/generate_qwen_edit.py deleted file mode 100644 index 3e7697c..0000000 --- a/src/mflux/generate_qwen_edit.py +++ /dev/null @@ -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() diff --git a/src/mflux/models/qwen/variants/edit/__init__.py b/src/mflux/models/qwen/variants/edit/__init__.py index 8833a8b..5180da1 100644 --- a/src/mflux/models/qwen/variants/edit/__init__.py +++ b/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"] diff --git a/src/mflux/models/qwen/variants/edit/qwen_image_edit.py b/src/mflux/models/qwen/variants/edit/qwen_image_edit.py deleted file mode 100644 index b453f55..0000000 --- a/src/mflux/models/qwen/variants/edit/qwen_image_edit.py +++ /dev/null @@ -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) diff --git a/tests/image_generation/test_generate_image_edit.py b/tests/image_generation/test_generate_image_edit.py index 47070fb..91cb06c 100644 --- a/tests/image_generation/test_generate_image_edit.py +++ b/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", diff --git a/tests/resources/reference_qwen_edit.png b/tests/resources/reference_qwen_edit.png deleted file mode 100644 index 3f9bb45..0000000 Binary files a/tests/resources/reference_qwen_edit.png and /dev/null differ