initial commi

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,6 +1,6 @@
 ---
 title: DreamOmni2 Gen
-emoji: 🌖
+emoji: 🖼️
 colorFrom: pink
 colorTo: purple
 sdk: gradio

app.py

@@ -0,0 +1,228 @@
+import os
+import re
+from PIL import Image
+import spaces
+import gradio as gr
+import uuid
+import argparse
+from huggingface_hub import login, snapshot_download
+
+import torch
+from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline
+from diffusers.utils import load_image
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from utils.vprocess import process_vision_info, resizeinput
+
+
+def extract_gen_content(text):
+    text = text[6:-7]
+    return text
+
+def _load_model_processor():
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    local_dir = snapshot_download(
+        repo_id="xiabs/DreamOmni2",
+        revision="main",
+        allow_patterns=["vlm-model/**", "gen_lora/**"],
+    )
+    vlm_dir = os.path.join(local_dir, 'vlm-model')
+    lora_dir = os.path.join(local_dir, 'gen_lora')
+
+    print(f"Loading models from vlm_path: {vlm_dir}, gen_lora_path: {lora_dir}")
+    pipe = DreamOmni2Pipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-Kontext-dev",
+        torch_dtype=torch.bfloat16
+    ).to(device)
+    pipe.load_lora_weights(lora_dir, adapter_name="generation")
+    pipe.set_adapters(["generation"], adapter_weights=[1])
+
+    vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        vlm_dir,
+        torch_dtype="bfloat16"
+    ).to(device)
+    processor = AutoProcessor.from_pretrained(vlm_dir)
+    return vlm_model, processor, pipe
+
+
+def _launch_demo(vlm_model, processor, pipe):
+
+    @spaces.GPU(duration=150)
+    def infer_vlm(input_img_path, input_instruction, prefix):
+        if not vlm_model or not processor:
+            raise gr.Error("VLM Model not loaded. Cannot process prompt.")
+        tp = []
+        for path in input_img_path:
+            tp.append({"type": "image", "image": path})
+        tp.append({"type": "text", "text": input_instruction + prefix})
+        messages = [{"role": "user", "content": tp}]
+
+        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
+        inputs = inputs.to(device=vlm_model.device)
+
+        generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
+        generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+        output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        return output_text[0]
+
+    PREFERRED_KONTEXT_RESOLUTIONS = [
+        (672, 1568),
+        (688, 1504),
+        (720, 1456),
+        (752, 1392),
+        (800, 1328),
+        (832, 1248),
+        (880, 1184),
+        (944, 1104),
+        (1024, 1024),
+        (1104, 944),
+        (1184, 880),
+        (1248, 832),
+        (1328, 800),
+        (1392, 752),
+        (1456, 720),
+        (1504, 688),
+        (1568, 672),
+    ]
+    def find_closest_resolution(width, height, preferred_resolutions):
+        input_ratio = width / height
+        closest_resolution = min(
+            preferred_resolutions,
+            key=lambda res: abs((res[0] / res[1]) - input_ratio)
+        )
+        return closest_resolution
+
+    @spaces.GPU(duration=150)
+    def perform_generation(input_img_paths, input_instruction, output_path, height=1024, width=1024):
+        prefix = " It is generation task."
+        source_imgs = []
+        for path in input_img_paths:
+            img = load_image(path)
+            # source_imgs.append(img)
+            source_imgs.append(resizeinput(img))
+        prompt = infer_vlm(input_img_paths, input_instruction, prefix)
+        prompt = extract_gen_content(prompt)
+        print(f"Generated Prompt for VLM: {prompt}")
+
+        image = pipe(
+            images=source_imgs,
+            height=height,
+            width=width,
+            prompt=prompt,
+            num_inference_steps=30,
+            guidance_scale=3.5,
+        ).images[0]
+        image.save(output_path)
+        print(f"Generation result saved to {output_path}")
+
+    @spaces.GPU(duration=150)
+    def process_request(image_file_1, image_file_2, instruction):
+        # debugpy.listen(5678)
+        # print("Waiting for debugger attach...")
+        # debugpy.wait_for_client()
+        if not image_file_1 or not image_file_2:
+            raise gr.Error("Please upload both images.")
+        if not instruction:
+            raise gr.Error("Please provide an instruction.")
+        if not pipe or not vlm_model:
+            raise gr.Error("Models not loaded. Check the console for errors.")
+        
+        output_path = f"/tmp/{uuid.uuid4()}.png"
+        input_img_paths = [image_file_1, image_file_2]  # List of file paths from the two gr.File inputs
+
+        perform_generation(input_img_paths, instruction, output_path)
+        return output_path
+
+    css = """
+    .text-center { text-align: center; }
+    .result-img img {
+        max-height: 60vh !important; 
+        min-height: 30vh !important;
+        width: auto !important;      
+        object-fit: contain;         
+    }
+    .input-img img {
+        max-height: 30vh !important; 
+        width: auto !important;      
+        object-fit: contain;         
+    }
+    """
+
+
+    with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo:
+        gr.HTML(
+            """
+            <h1 style="text-align:center; font-size:48px; font-weight:bold; margin-bottom:20px;">
+                DreamOmni2: Omni-purpose Image Generation and Editing
+            </h1>
+            """
+        )
+        gr.Markdown(
+            "Select a mode, upload two images, provide an instruction, and click 'Run'.",
+            elem_classes="text-center"
+        )
+        with gr.Row():
+            with gr.Column(scale=2):
+                gr.Markdown("⬆️ Upload images. Click or drag to upload.")
+                
+                with gr.Row():
+                    image_uploader_1 = gr.Image(
+                        label="Img 1",
+                        type="filepath",
+                        interactive=True,
+                        elem_classes="input-img",
+                    )
+                    image_uploader_2 = gr.Image(
+                        label="Img 2",
+                        type="filepath",
+                        interactive=True,
+                        elem_classes="input-img",
+                    )
+                
+                instruction_text = gr.Textbox(
+                    label="Instruction",
+                    lines=2,
+                    placeholder="Input your instruction for generation or editing here...",
+                )
+                run_button = gr.Button("Run", variant="primary")
+
+            with gr.Column(scale=2):
+                gr.Markdown("🖼️ **Generation Mode**: Create new scenes from reference images."
+                            "Tip: If the result is not what you expect, try clicking **Run** again. ")
+                output_image = gr.Image(
+                    label="Result",
+                    type="filepath",
+                    elem_classes="result-img",
+                )
+
+        # --- Examples ---
+        gr.Markdown("## Examples")
+
+        gr.Examples(
+            label="Generation Examples",
+            examples=[
+                [
+                    "example_input/gen_tests/img1.jpg",
+                    "example_input/gen_tests/img2.jpg",
+                    "In the scene, the character from the first image stands on the left, and the character from the second image stands on the right. They are shaking hands against the backdrop of a spaceship interior.",
+                    "example_input/gen_tests/gen_res.png"
+                ]
+            ],
+            inputs=[image_uploader_1, image_uploader_2, instruction_text, output_image],
+            cache_examples=False,
+        )
+
+        run_button.click(
+            fn=process_request,
+            inputs=[image_uploader_1, image_uploader_2, instruction_text],
+            outputs=output_image
+        )
+    
+    demo.launch()
+
+
+if __name__ == "__main__":
+    vlm_model, processor, pipe = _load_model_processor()
+    _launch_demo(vlm_model, processor, pipe)

dreamomni2/pipeline_dreamomni2.py

@@ -0,0 +1,1148 @@
+# Copyright 2025 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import FluxIPAdapterMixin, FluxLoraLoaderMixin, FromSingleFileMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, FluxTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    is_torch_xla_available,
+    logging,
+    replace_example_docstring,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import DreamOmni2Pipeline
+        >>> from diffusers.utils import load_image
+
+        >>> pipe = DreamOmni2Pipeline.from_pretrained(
+        ...     "black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16
+        ... )
+        >>> pipe.to("cuda")
+
+        >>> image = load_image(
+        ...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yarn-art-pikachu.png"
+        ... ).convert("RGB")
+        >>> prompt = "Make Pikachu hold a sign that says 'Black Forest Labs is awesome', yarn art style, detailed, vibrant colors"
+        >>> image = pipe(
+        ...     image=image,
+        ...     prompt=prompt,
+        ...     guidance_scale=2.5,
+        ...     generator=torch.Generator().manual_seed(42),
+        ... ).images[0]
+        >>> image.save("output.png")
+        ```
+"""
+
+PREFERRED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class DreamOmni2Pipeline(
+    DiffusionPipeline,
+    FluxLoraLoaderMixin,
+    FromSingleFileMixin,
+    TextualInversionLoaderMixin,
+    FluxIPAdapterMixin,
+):
+    r"""
+    The Flux Kontext pipeline for image-to-image and text-to-image generation.
+
+    Reference: https://bfl.ai/announcements/flux-1-kontext-dev
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+    """
+
+    model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->transformer->vae"
+    _optional_components = ["image_encoder", "feature_extractor"]
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_encoder=image_encoder,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.latent_channels = self.vae.config.latent_channels if getattr(self, "vae", None) else 16
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = (
+            self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
+        )
+        self.default_sample_size = 128
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_t5_prompt_embeds
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer_2)
+
+        text_inputs = self.tokenizer_2(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer_2(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer_2.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because `max_sequence_length` is set to "
+                f" {max_sequence_length} tokens: {removed_text}"
+            )
+
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+
+        dtype = self.text_encoder_2.dtype
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        _, seq_len, _ = prompt_embeds.shape
+
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._get_clip_prompt_embeds
+    def _get_clip_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]],
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+    ):
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt)
+
+        if isinstance(self, TextualInversionLoaderMixin):
+            prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer_max_length,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {self.tokenizer_max_length} tokens: {removed_text}"
+            )
+        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+        # Use pooled output of CLIPTextModel
+        prompt_embeds = prompt_embeds.pooler_output
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        prompt_2: Union[str, List[str]],
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        max_sequence_length: int = 512,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in all text-encoders
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        device = device or self._execution_device
+
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            if self.text_encoder is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder, lora_scale)
+            if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+                scale_lora_layers(self.text_encoder_2, lora_scale)
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        if prompt_embeds is None:
+            prompt_2 = prompt_2 or prompt
+            prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+            # We only use the pooled prompt output from the CLIPTextModel
+            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+                prompt=prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+            )
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+
+        if self.text_encoder is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder, lora_scale)
+
+        if self.text_encoder_2 is not None:
+            if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+                # Retrieve the original scale by scaling back the LoRA layers
+                unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+        dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+        return prompt_embeds, pooled_prompt_embeds, text_ids
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.encode_image
+    def encode_image(self, image, device, num_images_per_prompt):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if not isinstance(image, torch.Tensor):
+            image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+        image = image.to(device=device, dtype=dtype)
+        image_embeds = self.image_encoder(image).image_embeds
+        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        return image_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.prepare_ip_adapter_image_embeds
+    def prepare_ip_adapter_image_embeds(
+        self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt
+    ):
+        image_embeds = []
+        if ip_adapter_image_embeds is None:
+            if not isinstance(ip_adapter_image, list):
+                ip_adapter_image = [ip_adapter_image]
+
+            if len(ip_adapter_image) != self.transformer.encoder_hid_proj.num_ip_adapters:
+                raise ValueError(
+                    f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters."
+                )
+
+            for single_ip_adapter_image in ip_adapter_image:
+                single_image_embeds = self.encode_image(single_ip_adapter_image, device, 1)
+                image_embeds.append(single_image_embeds[None, :])
+        else:
+            if not isinstance(ip_adapter_image_embeds, list):
+                ip_adapter_image_embeds = [ip_adapter_image_embeds]
+
+            if len(ip_adapter_image_embeds) != self.transformer.encoder_hid_proj.num_ip_adapters:
+                raise ValueError(
+                    f"`ip_adapter_image_embeds` must have same length as the number of IP Adapters. Got {len(ip_adapter_image_embeds)} image embeds and {self.transformer.encoder_hid_proj.num_ip_adapters} IP Adapters."
+                )
+
+            for single_image_embeds in ip_adapter_image_embeds:
+                image_embeds.append(single_image_embeds)
+
+        ip_adapter_image_embeds = []
+        for single_image_embeds in image_embeds:
+            single_image_embeds = torch.cat([single_image_embeds] * num_images_per_prompt, dim=0)
+            single_image_embeds = single_image_embeds.to(device=device)
+            ip_adapter_image_embeds.append(single_image_embeds)
+
+        return ip_adapter_image_embeds
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        prompt_2,
+        height,
+        width,
+        negative_prompt=None,
+        negative_prompt_2=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        negative_pooled_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt_2 is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+            raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+        elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed. Make sure to generate `pooled_prompt_embeds` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed. Make sure to generate `negative_pooled_prompt_embeds` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 512:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 512 but is {max_sequence_length}")
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._prepare_latent_image_ids
+    def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+        latent_image_ids = torch.zeros(height, width, 3)
+        latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+        latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+        latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+        latent_image_ids = latent_image_ids.reshape(
+            latent_image_id_height * latent_image_id_width, latent_image_id_channels
+        )
+
+        return latent_image_ids.to(device=device, dtype=dtype)
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._pack_latents
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline._unpack_latents
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+        return latents
+
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
+
+        image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        return image_latents
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_slicing
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_slicing
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.enable_vae_tiling
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    # Copied from diffusers.pipelines.flux.pipeline_flux.FluxPipeline.disable_vae_tiling
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def prepare_latents(
+        self,
+        images: Optional[torch.Tensor],
+        batch_size: int,
+        num_channels_latents: int,
+        height: int,
+        width: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+    ):
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+        shape = (batch_size, num_channels_latents, height, width)
+        h_offset = 0
+        w_offset = 0
+        image_latents = image_ids = None
+        if images is not None:
+            tp_image_latents = []
+            tp_image_ids = []
+            for i, image in enumerate(images):
+                image = image.to(device=device, dtype=dtype)
+                if image.shape[1] != self.latent_channels:
+                    image_latents = self._encode_vae_image(image=image, generator=generator)
+                else:
+                    image_latents = image
+                if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+                    # expand init_latents for batch_size
+                    additional_image_per_prompt = batch_size // image_latents.shape[0]
+                    image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+                elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+                    raise ValueError(
+                        f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+                    )
+                else:
+                    image_latents = torch.cat([image_latents], dim=0)
+
+                image_latent_height, image_latent_width = image_latents.shape[2:]
+                image_latents = self._pack_latents(
+                    image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
+                )
+                image_ids = self._prepare_latent_image_ids(
+                    batch_size, image_latent_height // 2, image_latent_width // 2, device, dtype
+                )
+                image_ids[..., 0] = i+1
+                image_ids[..., 2] += w_offset
+                tp_image_latents.append(image_latents)
+                tp_image_ids.append(image_ids)
+                h_offset += image_latent_height //2
+                w_offset += image_latent_width //2
+            image_latents = torch.cat(tp_image_latents, dim=1)
+            image_ids = torch.cat(tp_image_ids, dim=0)
+
+        latent_ids = self._prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        return latents, image_latents, latent_ids, image_ids
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def joint_attention_kwargs(self):
+        return self._joint_attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        images: Optional[List[PipelineImageInput]] = None,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        true_cfg_scale: float = 1.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        max_area: int = 1024**2,
+        _auto_resize: bool = True,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
+                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
+                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
+                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
+                latents as `image`, but if passing latents directly it is not encoded again.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in all the text-encoders.
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                When > 1.0 and a provided `negative_prompt`, enables true classifier-free guidance.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 3.5):
+                Guidance scale as defined in [Classifier-Free Diffusion
+                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
+                the text `prompt`, usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*):
+                Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_ip_adapter_image:
+                (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512):
+                Maximum sequence length to use with the `prompt`.
+            max_area (`int`, defaults to `1024 ** 2`):
+                The maximum area of the generated image in pixels. The height and width will be adjusted to fit this
+                area while maintaining the aspect ratio.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_height, original_width = height, width
+        aspect_ratio = width / height
+        width = round((max_area * aspect_ratio) ** 0.5)
+        height = round((max_area / aspect_ratio) ** 0.5)
+
+        multiple_of = self.vae_scale_factor * 2
+        width = width // multiple_of * multiple_of
+        height = height // multiple_of * multiple_of
+
+        if height != original_height or width != original_width:
+            logger.warning(
+                f"Generation `height` and `width` have been adjusted to {height} and {width} to fit the model requirements."
+            )
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
+        )
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        if do_true_cfg:
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                negative_text_ids,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=lora_scale,
+            )
+
+        # 3. Preprocess image
+        if images is not None and not (isinstance(images[0], torch.Tensor) and images[0].size(1) == self.latent_channels):
+            tp_images=[]
+            for img in images:
+                image = img
+                image_height, image_width = self.image_processor.get_default_height_width(img)
+                aspect_ratio = image_width / image_height
+                if _auto_resize:
+                    # Kontext is trained on specific resolutions, using one of them is recommended
+                    _, image_width, image_height = min(
+                        (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
+                    )
+                image_width = image_width // multiple_of * multiple_of
+                image_height = image_height // multiple_of * multiple_of
+                image = self.image_processor.resize(image, image_height, image_width)
+                image = self.image_processor.preprocess(image, image_height, image_width)
+                tp_images.append(image)
+            images = tp_images
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, image_latents, latent_ids, image_ids = self.prepare_latents(
+            images,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        if image_ids is not None:
+            latent_ids = torch.cat([latent_ids, image_ids], dim=0)  # dim 0 is sequence dimension
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+            negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+        ):
+            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+            negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+        ):
+            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+            ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        negative_image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+                negative_ip_adapter_image,
+                negative_ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        # 6. Denoising loop
+        # We set the index here to remove DtoH sync, helpful especially during compilation.
+        # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred[:, : latents.size(1)]
+
+                if do_true_cfg:
+                    if negative_image_embeds is not None:
+                        self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=negative_text_ids,
+                        img_ids=latent_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)

inference_edit.py

@@ -0,0 +1,180 @@
+import torch
+try:
+    import torch_npu
+    from torch_npu.contrib import transfer_to_npu
+    import importlib
+    import transformers.utils
+    import transformers.models
+    origin_utils = transformers.utils
+    origin_models = transformers.models
+    import flash_attn
+    flash_attn.hack_transformers_flash_attn_2_available_check()
+    importlib.reload(transformers.utils)
+    importlib.reload(transformers.models)
+    origin_func = torch.nn.functional.interpolate
+    def new_func(input, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False):
+        if mode == "bilinear":
+            dtype = input.dtype
+            res = origin_func(input.to(torch.bfloat16), size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
+            return res.to(dtype)
+        else:
+            return origin_func(input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
+    torch.nn.functional.interpolate = new_func
+    from utils import patch_npu_record_stream
+    from utils import patch_npu_diffusers_get_1d_rotary_pos_embed
+    patch_npu_record_stream()
+    patch_npu_diffusers_get_1d_rotary_pos_embed()
+    USE_NPU = True
+except:
+    USE_NPU = False
+from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline
+from diffusers.utils import load_image
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+# from qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
+import os
+import argparse
+from tqdm import tqdm
+import json
+from PIL import Image
+import re
+import argparse 
+
+if USE_NPU:
+    device = "npu"
+else:
+    device = "cuda"
+
+def extract_gen_content(text):
+    text = text[6:-7]
+    
+    return text
+
+def parse_args():
+    """Parses command-line arguments for model paths and server configuration."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--vlm_path", 
+        type=str, 
+        default="./models/vlm-model", 
+        help="Path to the VLM model directory."
+    )
+    parser.add_argument(
+        "--edit_lora_path", 
+        type=str, 
+        default="./models/edit_lora", 
+        help="Path to the FLUX.1-Kontext editing LoRA weights directory."
+    )
+    parser.add_argument(
+        "--base_model_path", 
+        type=str, 
+        default="black-forest-labs/FLUX.1-Kontext-dev", 
+        help="Path to the FLUX.1-Kontext editing."
+    )
+    parser.add_argument(
+        "--input_img_path", 
+        type=str, 
+        nargs='+',  # Accept one or more input paths
+        default=["example_input/edit_tests/src.jpg", "example_input/edit_tests/ref.jpg"], 
+        help="List of input image paths (e.g., src and ref images)."
+    )
+    # Argument for the input instruction
+    parser.add_argument(
+        "--input_instruction", 
+        type=str, 
+        default="Make the woman from the second image stand on the road in the first image.", 
+        help="Instruction for image editing."
+    )
+    # Argument for the output image path
+    parser.add_argument(
+        "--output_path", 
+        type=str, 
+        default="example_input/edit_tests/edi_res.png", 
+        help="Path to save the output image."
+    )
+    
+    args = parser.parse_args()
+    return args
+
+ARGS = parse_args()
+vlm_path = ARGS.vlm_path
+edit_lora_path = ARGS.edit_lora_path
+base_model = ARGS.base_model_path
+pipe = DreamOmni2Pipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+pipe.to(device)
+
+pipe.load_lora_weights(
+    edit_lora_path,
+    adapter_name="edit" 
+)
+pipe.set_adapters(["edit"], adapter_weights=[1])   
+
+
+vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    vlm_path, torch_dtype="bfloat16", device_map="cuda"
+)
+processor = AutoProcessor.from_pretrained(vlm_path)
+
+def infer_vlm(input_img_path,input_instruction,prefix):
+    tp=[]
+    for path in input_img_path:
+        tp.append({"type": "image", "image": path})
+    tp.append({"type": "text", "text": input_instruction+prefix})
+    messages = [
+            {
+                "role": "user",
+                "content": tp,
+            }
+        ]
+
+    # Preparation for inference
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
+
+    # Inference
+    generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return output_text[0]
+
+def infer(source_imgs,prompt):
+    image = pipe(
+    images=source_imgs,
+    height=source_imgs[0].height,
+    width=source_imgs[0].width,
+    prompt=prompt,
+    num_inference_steps=30,
+    guidance_scale=3.5,
+    ).images[0]
+    return image
+    
+
+input_img_path=ARGS.input_img_path
+input_instruction=ARGS.input_instruction
+
+prefix=" It is editing task."
+source_imgs = []
+for path in input_img_path:
+    img = load_image(path)
+    # source_imgs.append(img)
+    source_imgs.append(resizeinput(img))
+
+prompt=infer_vlm(input_img_path,input_instruction,prefix)
+prompt = extract_gen_content(prompt)
+image=infer(source_imgs,prompt)
+output_path = ARGS.output_path
+image.save(output_path)

inference_gen.py

@@ -0,0 +1,192 @@
+import torch
+try:
+    import torch_npu
+    from torch_npu.contrib import transfer_to_npu
+    import importlib
+    import transformers.utils
+    import transformers.models
+    origin_utils = transformers.utils
+    origin_models = transformers.models
+    import flash_attn
+    flash_attn.hack_transformers_flash_attn_2_available_check()
+    importlib.reload(transformers.utils)
+    importlib.reload(transformers.models)
+    origin_func = torch.nn.functional.interpolate
+    def new_func(input, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False):
+        if mode == "bilinear":
+            dtype = input.dtype
+            res = origin_func(input.to(torch.bfloat16), size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
+            return res.to(dtype)
+        else:
+            return origin_func(input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias)
+    torch.nn.functional.interpolate = new_func
+    from utils import patch_npu_record_stream
+    from utils import patch_npu_diffusers_get_1d_rotary_pos_embed
+    patch_npu_record_stream()
+    patch_npu_diffusers_get_1d_rotary_pos_embed()
+    USE_NPU = True
+except:
+    USE_NPU = False
+from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline
+from diffusers.utils import load_image
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+# from qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
+import os
+import argparse
+from tqdm import tqdm
+import json
+from PIL import Image
+import re
+import argparse 
+
+if USE_NPU:
+    device = "npu"
+else:
+    device = "cuda"
+
+def extract_gen_content(text):
+    text = text[6:-7]
+    
+    return text
+
+def parse_args():
+    """Parses command-line arguments for model paths and server configuration."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--vlm_path", 
+        type=str, 
+        default="./models/vlm-model", 
+        help="Path to the VLM model directory."
+    )
+    parser.add_argument(
+        "--gen_lora_path", 
+        type=str, 
+        default="./models/gen_lora", 
+        help="Path to the FLUX.1-Kontext generation LoRA weights directory."
+    )
+    parser.add_argument(
+        "--base_model_path", 
+        type=str, 
+        default="black-forest-labs/FLUX.1-Kontext-dev", 
+        help="Path to the FLUX.1-Kontext editing."
+    )
+    parser.add_argument(
+        "--input_img_path", 
+        type=str, 
+        nargs='+',  # Accept one or more input paths
+        default=["example_input/gen_tests/img1.jpg","example_input/gen_tests/img2.jpg"], 
+        help="List of input image paths (e.g., src and ref images)."
+    )
+    # Argument for the input instruction
+    parser.add_argument(
+        "--input_instruction", 
+        type=str, 
+        default="In the scene, the character from the first image stands on the left, and the character from the second image stands on the right. They are shaking hands against the backdrop of a spaceship interior.", 
+        help="Instruction for image generation."
+    )
+    parser.add_argument(
+        "--height", 
+        type=int, 
+        default=1024, 
+        help="The height of output image."
+    )
+    parser.add_argument(
+        "--width", 
+        type=int, 
+        default=1024, 
+        help="The width of output image."
+    )
+    # Argument for the output image path
+    parser.add_argument(
+        "--output_path", 
+        type=str, 
+        default="example_input/gen_tests/gen_res.png", 
+        help="Path to save the output image."
+    )
+    
+    args = parser.parse_args()
+    return args
+
+ARGS = parse_args()
+vlm_path = ARGS.vlm_path
+gen_lora_path = ARGS.gen_lora_path
+base_model = ARGS.base_model_path
+pipe = DreamOmni2Pipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+pipe.to(device)
+
+pipe.load_lora_weights(
+    gen_lora_path,
+    adapter_name="generation"  
+)
+pipe.set_adapters(["generation"], adapter_weights=[1])   
+
+
+vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    vlm_path, torch_dtype="bfloat16", device_map="cuda"
+)
+processor = AutoProcessor.from_pretrained(vlm_path)
+
+def infer_vlm(input_img_path,input_instruction,prefix):
+    tp=[]
+    for path in input_img_path:
+        tp.append({"type": "image", "image": path})
+    tp.append({"type": "text", "text": input_instruction+prefix})
+    messages = [
+            {
+                "role": "user",
+                "content": tp,
+            }
+        ]
+
+    # Preparation for inference
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
+
+    # Inference
+    generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return output_text[0]
+
+def infer(source_imgs,prompt,height=1024,width=1024):
+    image = pipe(
+    images=source_imgs,
+    height=height,
+    width=width,
+    prompt=prompt,
+    num_inference_steps=30,
+    guidance_scale=3.5,
+    ).images[0]
+    return image
+    
+
+input_img_path=ARGS.input_img_path
+input_instruction=ARGS.input_instruction
+
+prefix=" It is generation task."
+source_imgs = []
+for path in input_img_path:
+    img = load_image(path)
+    # source_imgs.append(img)
+    source_imgs.append(resizeinput(img))
+
+prompt=infer_vlm(input_img_path,input_instruction,prefix)
+prompt = extract_gen_content(prompt)
+image=infer(source_imgs,prompt,height=ARGS.height,width=ARGS.width)
+output_path = ARGS.output_path
+image.save(output_path)

my_datasets/.gitkeep

@@ -0,0 +1,327 @@
+import os
+import re
+import json
+import gc
+import functools
+import contextlib
+from typing import Dict, Union, Optional, Type, Set
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
+    StateDictType,
+    FullOptimStateDictConfig,
+    FullStateDictConfig,
+)
+import torch.distributed.checkpoint as torch_dcp
+import torch.distributed.checkpoint.state_dict
+from torch.distributed.fsdp.api import (
+    ShardingStrategy,
+    BackwardPrefetch,
+    MixedPrecision,
+)
+import accelerate
+import safetensors
+import diffusers
+import transformers
+from huggingface_hub.serialization import split_torch_state_dict_into_shards
+import os, re, json
+from typing import Union
+import torch
+import safetensors.torch
+import accelerate
+# from .ema_utils import EMAModel
+
+
+def upcast_trainable_param_to_fp32_(fsdp_model):
+    for m in FSDP.fsdp_modules(fsdp_model):
+        if m._has_params:
+            param = m._flat_param
+            if (
+                param.dtype != torch.float32
+                and param.device != torch.device("meta")
+                and param.requires_grad
+            ):
+                param.data = param.data.to(torch.float32)
+                m._handle._orig_param_dtype = torch.float32
+
+
+def get_module_to_ignore_mixed_precision():
+    try:
+        from apex.normalization import FusedLayerNorm
+
+        return [
+            torch.nn.GroupNorm,
+            torch.nn.modules.batchnorm._BatchNorm,
+            torch.nn.LayerNorm,
+            FusedLayerNorm,
+        ]
+    except:
+        return [
+            torch.nn.GroupNorm,
+            torch.nn.modules.batchnorm._BatchNorm,
+            torch.nn.LayerNorm,
+        ]
+
+
+def is_fsdp_model(model):
+    return len(FSDP.fsdp_modules(model)) > 0
+
+
+def size_based_auto_wrap_policy(
+    module: torch.nn.Module,
+    recurse: bool,
+    nonwrapped_numel: int,
+    # Additional custom arguments
+    min_num_params: int = int(1e8),
+    force_leaf_modules: Optional[Set[Type[torch.nn.Module]]] = None,
+    exclude_wrap_modules: Optional[Set[Type[torch.nn.Module]]] = None,
+) -> bool:
+    """
+    A size-based auto wrap policy.
+
+    Args:
+        module (nn.Module): Current module being considered.
+        recurse (bool): If ``False``, then this function must decide whether
+            ``module`` should be wrapped as an FSDP instance or not. If
+            ``True``, then the function is still recursing down the module
+            tree as a part of the DFS.
+        nonwrapped_numel (int): Parameter numel not yet wrapped.
+
+        min_num_params (int): Customizable policy input that controls the size
+            threshold over which a module is ready to be wrapped. This is in
+            units of numel.
+        force_leaf_modules (Set[Type[nn.Module]]): Set of module types to keep
+            as leaves, i.e. their children will never be wrapped.
+        exclude_wrap_modules (Set[Type[nn.Module]]): Set of module types to be
+            excluded in wrapping.
+
+    Returns:
+        Whether ``module`` should be wrapped.
+    """
+    force_leaf_modules = (
+        size_based_auto_wrap_policy.FORCE_LEAF_MODULES  # type: ignore[attr-defined]
+        if force_leaf_modules is None
+        else force_leaf_modules
+    )
+    exclude_wrap_modules = (
+        size_based_auto_wrap_policy.EXCLUDE_WRAP_MODULES  # type: ignore[attr-defined]
+        if exclude_wrap_modules is None
+        else exclude_wrap_modules
+    )
+
+    # Keep the argument `min_num_params` for BC for now, but it represents the
+    # minimum non-wrapped *numel* before triggering a wrapping
+    min_nonwrapped_numel = min_num_params
+    is_large = nonwrapped_numel >= min_nonwrapped_numel
+    STOP_FLAG_NAME = "__FSDP_STOP_WARP_FLAG_CUSTOM_POLICY_size_based_auto_wrap_policy"
+    if recurse:
+        # use MixedPrecision cause ALWAYS recurse
+        if isinstance(module, tuple(force_leaf_modules)):
+            for m in module.children():
+                m.apply(lambda m: setattr(m, STOP_FLAG_NAME, True))
+        return True
+    else:
+        if getattr(module, size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME, False):
+            return True
+        elif getattr(module, STOP_FLAG_NAME, False):
+            return False
+        else:
+            # If we are not recursing, determine if we should wrap.
+            return is_large and not isinstance(module, tuple(exclude_wrap_modules))
+
+
+# Set those defaults to the size_based_auto_wrap_policy function. Make them easy to be imported.
+size_based_auto_wrap_policy.EXCLUDE_WRAP_MODULES = {torch.nn.ModuleList, torch.nn.ModuleDict}  # type: ignore[attr-defined]
+size_based_auto_wrap_policy.FORCE_LEAF_MODULES = {torch.nn.MultiheadAttention}  # type: ignore[attr-defined]
+size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME = (
+    "__FSDP_LEAF_ROOT_FLAG_CUSTOM_POLICY_size_based_auto_wrap_policy"
+)
+
+
+def mark_leaf_root_(module):
+    setattr(
+        module,
+        size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME,
+        True,
+    )
+
+
+def make_model_fsdp(
+    model,
+    param_dtype,
+    device,
+    reduce_dtype=None,
+    buffer_dtype=None,
+    sync_module_states=True,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+    module_classes_to_ignore_mixed_precision=None,
+    ignored_states=None,
+    ignored_modules=None,
+    auto_wrap_policy=None,
+    part_size=1e6,
+    force_leaf_modules=None,
+    exclude_wrap_modules=None,
+    use_orig_params=False
+):
+    if module_classes_to_ignore_mixed_precision is None:
+        module_classes_to_ignore_mixed_precision = (
+            get_module_to_ignore_mixed_precision()
+        )
+    if auto_wrap_policy is not None:
+        auto_wrap_policy = auto_wrap_policy
+    elif sharding_strategy == ShardingStrategy.NO_SHARD:
+        auto_wrap_policy = None
+    else:
+        auto_wrap_policy = functools.partial(
+            size_based_auto_wrap_policy,
+            min_num_params=part_size,
+            force_leaf_modules=force_leaf_modules,
+            exclude_wrap_modules=exclude_wrap_modules,
+        )
+
+    model = FSDP(
+        model,
+        sharding_strategy=sharding_strategy,
+        process_group=process_group,
+        forward_prefetch=True,
+        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+        limit_all_gathers=True,
+        use_orig_params=use_orig_params,
+        sync_module_states=sync_module_states,
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype or torch.float32,
+            buffer_dtype=buffer_dtype or torch.float32,
+            keep_low_precision_grads=False,
+            cast_forward_inputs=False,
+            cast_root_forward_inputs=True,
+            _module_classes_to_ignore=module_classes_to_ignore_mixed_precision,
+        ),
+        auto_wrap_policy=auto_wrap_policy,
+        ignored_states=ignored_states,
+        ignored_modules=ignored_modules,
+        device_id=device,
+    )
+    torch.cuda.empty_cache()
+    gc.collect()
+    return model
+
+
+def save_fsdp_lora(
+    model_to_save,                      # FSDP 包裹的模型
+    save_directory: Union[str, os.PathLike],
+    is_main_process: bool = True,
+    lora_regex: str = r"(?:lora)",  # 根据自己命名习惯调
+):
+    """
+    仅保存 LoRA 层的权重。适用于 FSDP 并与 safetensors 兼容。
+    """
+    # 1. 解包 FSDP，拿到裸模型
+    unwrapped_model = accelerate.utils.extract_model_from_parallel(model_to_save)
+
+    # 2. 创建保存目录
+    if is_main_process:
+        os.makedirs(save_directory, exist_ok=True)
+
+    # 3. 收集完整 state_dict（CPU 上）
+    state_dict = torch_dcp.state_dict.get_model_state_dict(
+        model_to_save,
+        options=torch_dcp.state_dict.StateDictOptions(
+            full_state_dict=True,
+            cpu_offload=True,
+            ignore_frozen_params=False,
+        ),
+    )
+
+    # 4. 过滤出 LoRA 参数
+    lora_pattern = re.compile(lora_regex)
+    lora_state_dict = {
+        k: v for k, v in state_dict.items() if lora_pattern.search(k) is not None
+    }
+
+    if not lora_state_dict:
+        raise ValueError(
+            "未找到匹配 LoRA 的参数。请检查 lora_regex 是否符合命名规则。"
+        )
+
+    # 5. 保存为单文件 *.safetensors
+    if is_main_process:
+        weight_file = os.path.join(save_directory, "adapter_model.safetensors")
+        safetensors.torch.save_file(
+            lora_state_dict, weight_file, metadata={"format": "pt", "type": "lora"}
+        )
+
+
+def load_fsdp_model_(model_to_load: FSDP, save_directory: Union[str, os.PathLike]):
+    with FSDP.state_dict_type(
+        model_to_load,
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        state_dict_config=FullStateDictConfig(
+            rank0_only=False,
+        ),
+    ):
+        _model = model_to_load.from_pretrained(save_directory)
+        model_to_load.load_state_dict(_model.state_dict())
+
+
+def save_fsdp_optimizer(
+    models: Dict,
+    optimizer_to_save: torch.optim.Optimizer,
+    save_directory: Union[str, os.PathLike],
+    is_main_process: bool = True,
+):
+    _fsdp_state_dict_config = dict(
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        optim_state_dict_config=FullOptimStateDictConfig(
+            offload_to_cpu=True,
+            rank0_only=True,
+        ),
+    )
+    mgrs = list()
+    for m in models.values():
+        if len(FSDP.fsdp_modules(m)) > 0:
+            mgrs.append(FSDP.state_dict_type(m, **_fsdp_state_dict_config))
+
+    with contextlib.ExitStack() as stack:
+        for mgr in mgrs:
+            stack.enter_context(mgr)
+        optim_state_dict = FSDP.optim_state_dict(
+            torch.nn.ModuleDict(models),
+            optimizer_to_save,
+        )
+        if is_main_process:
+            torch.save(
+                optim_state_dict, os.path.join(save_directory, "optim_states.pth")
+            )
+
+
+def load_fsdp_optimizer_(
+    models: Dict,
+    optimizer_to_load: torch.optim.Optimizer,
+    save_directory: Union[str, os.PathLike],
+):
+    _fsdp_state_dict_config = dict(
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        optim_state_dict_config=FullOptimStateDictConfig(
+            rank0_only=False,
+        ),
+    )
+    mgrs = list()
+    for m in models.values():
+        if len(FSDP.fsdp_modules(m)) > 0:
+            mgrs.append(FSDP.state_dict_type(m, **_fsdp_state_dict_config))
+
+    with contextlib.ExitStack() as stack:
+        for mgr in mgrs:
+            stack.enter_context(mgr)
+        optimizer_path = os.path.join(save_directory, "optim_states.pth")
+        assert os.path.isfile(optimizer_path)
+        optim_state_dict = torch.load(optimizer_path)
+        optim_state_dict = FSDP.optim_state_dict_to_load(
+            torch.nn.ModuleDict(models),
+            optimizer_to_load,
+            optim_state_dict,
+        )
+        optimizer_to_load.load_state_dict(optim_state_dict)

requirements.txt

@@ -0,0 +1,16 @@
+timm==1.0.20
+ujson
+peft==0.17.1
+datasets==4.1.1
+transformers==4.56.2
+opencv-python
+qwen-vl-utils==0.0.14
+lmdb==1.7.3
+diffusers==0.35.1
+numpy==1.26.4
+torch==2.8.0
+torchaudio==2.8.0
+torchvision==0.23.0
+gradio==5.47.2
+sentencepiece
+safetensors

utils/fsdp_utils.py

@@ -0,0 +1,327 @@
+import os
+import re
+import json
+import gc
+import functools
+import contextlib
+from typing import Dict, Union, Optional, Type, Set
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
+    StateDictType,
+    FullOptimStateDictConfig,
+    FullStateDictConfig,
+)
+import torch.distributed.checkpoint as torch_dcp
+import torch.distributed.checkpoint.state_dict
+from torch.distributed.fsdp.api import (
+    ShardingStrategy,
+    BackwardPrefetch,
+    MixedPrecision,
+)
+import accelerate
+import safetensors
+import diffusers
+import transformers
+from huggingface_hub.serialization import split_torch_state_dict_into_shards
+import os, re, json
+from typing import Union
+import torch
+import safetensors.torch
+import accelerate
+# from .ema_utils import EMAModel
+
+
+def upcast_trainable_param_to_fp32_(fsdp_model):
+    for m in FSDP.fsdp_modules(fsdp_model):
+        if m._has_params:
+            param = m._flat_param
+            if (
+                param.dtype != torch.float32
+                and param.device != torch.device("meta")
+                and param.requires_grad
+            ):
+                param.data = param.data.to(torch.float32)
+                m._handle._orig_param_dtype = torch.float32
+
+
+def get_module_to_ignore_mixed_precision():
+    try:
+        from apex.normalization import FusedLayerNorm
+
+        return [
+            torch.nn.GroupNorm,
+            torch.nn.modules.batchnorm._BatchNorm,
+            torch.nn.LayerNorm,
+            FusedLayerNorm,
+        ]
+    except:
+        return [
+            torch.nn.GroupNorm,
+            torch.nn.modules.batchnorm._BatchNorm,
+            torch.nn.LayerNorm,
+        ]
+
+
+def is_fsdp_model(model):
+    return len(FSDP.fsdp_modules(model)) > 0
+
+
+def size_based_auto_wrap_policy(
+    module: torch.nn.Module,
+    recurse: bool,
+    nonwrapped_numel: int,
+    # Additional custom arguments
+    min_num_params: int = int(1e8),
+    force_leaf_modules: Optional[Set[Type[torch.nn.Module]]] = None,
+    exclude_wrap_modules: Optional[Set[Type[torch.nn.Module]]] = None,
+) -> bool:
+    """
+    A size-based auto wrap policy.
+
+    Args:
+        module (nn.Module): Current module being considered.
+        recurse (bool): If ``False``, then this function must decide whether
+            ``module`` should be wrapped as an FSDP instance or not. If
+            ``True``, then the function is still recursing down the module
+            tree as a part of the DFS.
+        nonwrapped_numel (int): Parameter numel not yet wrapped.
+
+        min_num_params (int): Customizable policy input that controls the size
+            threshold over which a module is ready to be wrapped. This is in
+            units of numel.
+        force_leaf_modules (Set[Type[nn.Module]]): Set of module types to keep
+            as leaves, i.e. their children will never be wrapped.
+        exclude_wrap_modules (Set[Type[nn.Module]]): Set of module types to be
+            excluded in wrapping.
+
+    Returns:
+        Whether ``module`` should be wrapped.
+    """
+    force_leaf_modules = (
+        size_based_auto_wrap_policy.FORCE_LEAF_MODULES  # type: ignore[attr-defined]
+        if force_leaf_modules is None
+        else force_leaf_modules
+    )
+    exclude_wrap_modules = (
+        size_based_auto_wrap_policy.EXCLUDE_WRAP_MODULES  # type: ignore[attr-defined]
+        if exclude_wrap_modules is None
+        else exclude_wrap_modules
+    )
+
+    # Keep the argument `min_num_params` for BC for now, but it represents the
+    # minimum non-wrapped *numel* before triggering a wrapping
+    min_nonwrapped_numel = min_num_params
+    is_large = nonwrapped_numel >= min_nonwrapped_numel
+    STOP_FLAG_NAME = "__FSDP_STOP_WARP_FLAG_CUSTOM_POLICY_size_based_auto_wrap_policy"
+    if recurse:
+        # use MixedPrecision cause ALWAYS recurse
+        if isinstance(module, tuple(force_leaf_modules)):
+            for m in module.children():
+                m.apply(lambda m: setattr(m, STOP_FLAG_NAME, True))
+        return True
+    else:
+        if getattr(module, size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME, False):
+            return True
+        elif getattr(module, STOP_FLAG_NAME, False):
+            return False
+        else:
+            # If we are not recursing, determine if we should wrap.
+            return is_large and not isinstance(module, tuple(exclude_wrap_modules))
+
+
+# Set those defaults to the size_based_auto_wrap_policy function. Make them easy to be imported.
+size_based_auto_wrap_policy.EXCLUDE_WRAP_MODULES = {torch.nn.ModuleList, torch.nn.ModuleDict}  # type: ignore[attr-defined]
+size_based_auto_wrap_policy.FORCE_LEAF_MODULES = {torch.nn.MultiheadAttention}  # type: ignore[attr-defined]
+size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME = (
+    "__FSDP_LEAF_ROOT_FLAG_CUSTOM_POLICY_size_based_auto_wrap_policy"
+)
+
+
+def mark_leaf_root_(module):
+    setattr(
+        module,
+        size_based_auto_wrap_policy.LEAF_ROOT_FLAG_NAME,
+        True,
+    )
+
+
+def make_model_fsdp(
+    model,
+    param_dtype,
+    device,
+    reduce_dtype=None,
+    buffer_dtype=None,
+    sync_module_states=True,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.HYBRID_SHARD,
+    module_classes_to_ignore_mixed_precision=None,
+    ignored_states=None,
+    ignored_modules=None,
+    auto_wrap_policy=None,
+    part_size=1e6,
+    force_leaf_modules=None,
+    exclude_wrap_modules=None,
+    use_orig_params=False
+):
+    if module_classes_to_ignore_mixed_precision is None:
+        module_classes_to_ignore_mixed_precision = (
+            get_module_to_ignore_mixed_precision()
+        )
+    if auto_wrap_policy is not None:
+        auto_wrap_policy = auto_wrap_policy
+    elif sharding_strategy == ShardingStrategy.NO_SHARD:
+        auto_wrap_policy = None
+    else:
+        auto_wrap_policy = functools.partial(
+            size_based_auto_wrap_policy,
+            min_num_params=part_size,
+            force_leaf_modules=force_leaf_modules,
+            exclude_wrap_modules=exclude_wrap_modules,
+        )
+
+    model = FSDP(
+        model,
+        sharding_strategy=sharding_strategy,
+        process_group=process_group,
+        forward_prefetch=True,
+        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+        limit_all_gathers=True,
+        use_orig_params=use_orig_params,
+        sync_module_states=sync_module_states,
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype or torch.float32,
+            buffer_dtype=buffer_dtype or torch.float32,
+            keep_low_precision_grads=False,
+            cast_forward_inputs=False,
+            cast_root_forward_inputs=True,
+            _module_classes_to_ignore=module_classes_to_ignore_mixed_precision,
+        ),
+        auto_wrap_policy=auto_wrap_policy,
+        ignored_states=ignored_states,
+        ignored_modules=ignored_modules,
+        device_id=device,
+    )
+    torch.cuda.empty_cache()
+    gc.collect()
+    return model
+
+
+def save_fsdp_lora(
+    model_to_save,                      # FSDP 包裹的模型
+    save_directory: Union[str, os.PathLike],
+    is_main_process: bool = True,
+    lora_regex: str = r"(?:lora)",  # 根据自己命名习惯调
+):
+    """
+    仅保存 LoRA 层的权重。适用于 FSDP 并与 safetensors 兼容。
+    """
+    # 1. 解包 FSDP，拿到裸模型
+    unwrapped_model = accelerate.utils.extract_model_from_parallel(model_to_save)
+
+    # 2. 创建保存目录
+    if is_main_process:
+        os.makedirs(save_directory, exist_ok=True)
+
+    # 3. 收集完整 state_dict（CPU 上）
+    state_dict = torch_dcp.state_dict.get_model_state_dict(
+        model_to_save,
+        options=torch_dcp.state_dict.StateDictOptions(
+            full_state_dict=True,
+            cpu_offload=True,
+            ignore_frozen_params=False,
+        ),
+    )
+
+    # 4. 过滤出 LoRA 参数
+    lora_pattern = re.compile(lora_regex)
+    lora_state_dict = {
+        k: v for k, v in state_dict.items() if lora_pattern.search(k) is not None
+    }
+
+    if not lora_state_dict:
+        raise ValueError(
+            "未找到匹配 LoRA 的参数。请检查 lora_regex 是否符合命名规则。"
+        )
+
+    # 5. 保存为单文件 *.safetensors
+    if is_main_process:
+        weight_file = os.path.join(save_directory, "adapter_model.safetensors")
+        safetensors.torch.save_file(
+            lora_state_dict, weight_file, metadata={"format": "pt", "type": "lora"}
+        )
+
+
+def load_fsdp_model_(model_to_load: FSDP, save_directory: Union[str, os.PathLike]):
+    with FSDP.state_dict_type(
+        model_to_load,
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        state_dict_config=FullStateDictConfig(
+            rank0_only=False,
+        ),
+    ):
+        _model = model_to_load.from_pretrained(save_directory)
+        model_to_load.load_state_dict(_model.state_dict())
+
+
+def save_fsdp_optimizer(
+    models: Dict,
+    optimizer_to_save: torch.optim.Optimizer,
+    save_directory: Union[str, os.PathLike],
+    is_main_process: bool = True,
+):
+    _fsdp_state_dict_config = dict(
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        optim_state_dict_config=FullOptimStateDictConfig(
+            offload_to_cpu=True,
+            rank0_only=True,
+        ),
+    )
+    mgrs = list()
+    for m in models.values():
+        if len(FSDP.fsdp_modules(m)) > 0:
+            mgrs.append(FSDP.state_dict_type(m, **_fsdp_state_dict_config))
+
+    with contextlib.ExitStack() as stack:
+        for mgr in mgrs:
+            stack.enter_context(mgr)
+        optim_state_dict = FSDP.optim_state_dict(
+            torch.nn.ModuleDict(models),
+            optimizer_to_save,
+        )
+        if is_main_process:
+            torch.save(
+                optim_state_dict, os.path.join(save_directory, "optim_states.pth")
+            )
+
+
+def load_fsdp_optimizer_(
+    models: Dict,
+    optimizer_to_load: torch.optim.Optimizer,
+    save_directory: Union[str, os.PathLike],
+):
+    _fsdp_state_dict_config = dict(
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        optim_state_dict_config=FullOptimStateDictConfig(
+            rank0_only=False,
+        ),
+    )
+    mgrs = list()
+    for m in models.values():
+        if len(FSDP.fsdp_modules(m)) > 0:
+            mgrs.append(FSDP.state_dict_type(m, **_fsdp_state_dict_config))
+
+    with contextlib.ExitStack() as stack:
+        for mgr in mgrs:
+            stack.enter_context(mgr)
+        optimizer_path = os.path.join(save_directory, "optim_states.pth")
+        assert os.path.isfile(optimizer_path)
+        optim_state_dict = torch.load(optimizer_path)
+        optim_state_dict = FSDP.optim_state_dict_to_load(
+            torch.nn.ModuleDict(models),
+            optimizer_to_load,
+            optim_state_dict,
+        )
+        optimizer_to_load.load_state_dict(optim_state_dict)

utils/infer_utils.py

@@ -0,0 +1,163 @@
+import os
+import torch
+def tokenize_prompt(tokenizer, prompt, max_sequence_length):
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    return text_input_ids
+
+def _encode_prompt_with_t5(
+    text_encoder,
+    tokenizer,
+    max_sequence_length=512,
+    prompt=None,
+    num_images_per_prompt=1,
+    device=None,
+    text_input_ids=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device))[0]
+
+    if hasattr(text_encoder, "module"):
+        dtype = text_encoder.module.dtype
+    else:
+        dtype = text_encoder.dtype
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    _, seq_len, _ = prompt_embeds.shape
+
+    # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+    return prompt_embeds
+
+
+def _encode_prompt_with_clip(
+    text_encoder,
+    tokenizer,
+    prompt: str,
+    device=None,
+    text_input_ids=None,
+    num_images_per_prompt: int = 1,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+    if hasattr(text_encoder, "module"):
+        dtype = text_encoder.module.dtype
+    else:
+        dtype = text_encoder.dtype
+    # Use pooled output of CLIPTextModel
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    # duplicate text embeddings for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+    return prompt_embeds
+
+
+def encode_prompt(
+    text_encoders,
+    tokenizers,
+    prompt: str,
+    max_sequence_length,
+    device=None,
+    num_images_per_prompt: int = 1,
+    text_input_ids_list=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+
+    if hasattr(text_encoders[0], "module"):
+        dtype = text_encoders[0].module.dtype
+    else:
+        dtype = text_encoders[0].dtype
+
+    pooled_prompt_embeds = _encode_prompt_with_clip(
+        text_encoder=text_encoders[0],
+        tokenizer=tokenizers[0],
+        prompt=prompt,
+        device=device if device is not None else text_encoders[0].device,
+        num_images_per_prompt=num_images_per_prompt,
+        text_input_ids=text_input_ids_list[0] if text_input_ids_list else None,
+    )
+
+    prompt_embeds = _encode_prompt_with_t5(
+        text_encoder=text_encoders[1],
+        tokenizer=tokenizers[1],
+        max_sequence_length=max_sequence_length,
+        prompt=prompt,
+        num_images_per_prompt=num_images_per_prompt,
+        device=device if device is not None else text_encoders[1].device,
+        text_input_ids=text_input_ids_list[1] if text_input_ids_list else None,
+    )
+
+    text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+
+
+def compute_text_embeddings( args, accelerator, prompt, text_encoders, tokenizers):
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt(
+            text_encoders, tokenizers, prompt, args.max_sequence_length
+        )
+        prompt_embeds = prompt_embeds.to(accelerator.device)
+        pooled_prompt_embeds = pooled_prompt_embeds.to(accelerator.device)
+        text_ids = text_ids.to(accelerator.device)
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+
+def get_sigmas(noise_scheduler_copy,accelerator, timesteps, n_dim=4, dtype=torch.float32):
+    sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
+    schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
+    timesteps = timesteps.to(accelerator.device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+    sigma = sigmas[step_indices].flatten()
+    while len(sigma.shape) < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma

utils/init_utils.py

@@ -0,0 +1,48 @@
+from transformers import CLIPTokenizer, PretrainedConfig, T5TokenizerFast
+import logging
+def load_text_encoders(args, class_one, class_two):
+    text_encoder_one = class_one.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision, variant=args.variant
+    )
+    text_encoder_two = class_two.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder_2", revision=args.revision, variant=args.variant
+    )
+    return text_encoder_one, text_encoder_two
+
+def import_model_class_from_model_name_or_path(
+    pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder"
+):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path, subfolder=subfolder, revision=revision
+    )
+    model_class = text_encoder_config.architectures[0]
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "T5EncoderModel":
+        from transformers import T5EncoderModel
+
+        return T5EncoderModel
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+
+def create_logger(logging_dir,accelerator):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    if accelerator.is_main_process:  # real logger
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[\033[34m%(asctime)s\033[0m] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            handlers=[
+                logging.StreamHandler(),
+                logging.FileHandler(f"{logging_dir}/log.txt"),
+            ],
+        )
+        logger = logging.getLogger(__name__)
+    else:  # dummy logger (does nothing)
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger

utils/parser_config.py

@@ -0,0 +1,314 @@
+import argparse
+import os
+
+def parse_args(input_args=None):
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--revision",
+        type=str,
+        default=None,
+        required=False,
+        help="Revision of pretrained model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--vae_encode_mode",
+        type=str,
+        default="mode",
+        choices=["sample", "mode"],
+        help="VAE encoding mode.",
+    )
+    parser.add_argument(
+        "--variant",
+        type=str,
+        default=None,
+        help="Variant of the model files of the pretrained model identifier from huggingface.co/models, 'e.g.' fp16",
+    )
+    parser.add_argument("--repeats", type=int, default=1, help="How many times to repeat the training data.")
+    parser.add_argument(
+        "--max_sequence_length",
+        type=int,
+        default=512,
+        help="Maximum sequence length to use with with the T5 text encoder",
+    )
+    parser.add_argument(
+        "--rank",
+        type=int,
+        default=4,
+        help=("The dimension of the LoRA update matrices."),
+    )
+    parser.add_argument(
+        "--lora_alpha",
+        type=int,
+        default=4,
+        help="LoRA alpha to be used for additional scaling.",
+    )
+    parser.add_argument("--lora_dropout", type=float, default=0.0, help="Dropout probability for LoRA layers")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="flux-dreambooth-lora",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--aspect_ratio_buckets",
+        type=str,
+        default=None,
+        help=(
+            "Aspect ratio buckets to use for training. Define as a string of 'h1,w1;h2,w2;...'. "
+            "e.g. '1024,1024;768,1360;1360,768;880,1168;1168,880;1248,832;832,1248'"
+            "Images will be resized and cropped to fit the nearest bucket. If provided, --resolution is ignored."
+        ),
+    )
+    parser.add_argument(
+        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=1)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=(
+            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+            " training using `--resume_from_checkpoint`."
+        ),
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=(
+            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+        ),
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=3.5,
+        help="the FLUX.1 dev variant is a guidance distilled model",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of hard resets of the lr in cosine_with_restarts scheduler.",
+    )
+    parser.add_argument("--lr_power", type=float, default=1.0, help="Power factor of the polynomial scheduler.")
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=0,
+        help=(
+            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
+        ),
+    )
+    parser.add_argument(
+        "--weighting_scheme",
+        type=str,
+        default="none",
+        choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "none"],
+        help=('We default to the "none" weighting scheme for uniform sampling and uniform loss'),
+    )
+    parser.add_argument(
+        "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--mode_scale",
+        type=float,
+        default=1.29,
+        help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+    )
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        default="AdamW",
+        help=('The optimizer type to use. Choose between ["AdamW", "prodigy"]'),
+    )
+
+    parser.add_argument(
+        "--use_8bit_adam",
+        action="store_true",
+        help="Whether or not to use 8-bit Adam from bitsandbytes. Ignored if optimizer is not set to AdamW",
+    )
+
+    parser.add_argument(
+        "--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam and Prodigy optimizers."
+    )
+    parser.add_argument(
+        "--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam and Prodigy optimizers."
+    )
+    parser.add_argument(
+        "--prodigy_beta3",
+        type=float,
+        default=None,
+        help="coefficients for computing the Prodigy stepsize using running averages. If set to None, "
+        "uses the value of square root of beta2. Ignored if optimizer is adamW",
+    )
+    parser.add_argument("--prodigy_decouple", type=bool, default=True, help="Use AdamW style decoupled weight decay")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-04, help="Weight decay to use for unet params")
+    parser.add_argument(
+        "--adam_weight_decay_text_encoder", type=float, default=1e-03, help="Weight decay to use for text_encoder"
+    )
+
+    parser.add_argument(
+        "--lora_layers",
+        type=str,
+        default=None,
+        help=(
+            'The transformer modules to apply LoRA training on. Please specify the layers in a comma separated. E.g. - "to_k,to_q,to_v,to_out.0" will result in lora training of attention layers only'
+        ),
+    )
+
+    parser.add_argument(
+        "--adam_epsilon",
+        type=float,
+        default=1e-08,
+        help="Epsilon value for the Adam optimizer and Prodigy optimizers.",
+    )
+
+    parser.add_argument(
+        "--prodigy_use_bias_correction",
+        type=bool,
+        default=True,
+        help="Turn on Adam's bias correction. True by default. Ignored if optimizer is adamW",
+    )
+    parser.add_argument(
+        "--prodigy_safeguard_warmup",
+        type=bool,
+        default=True,
+        help="Remove lr from the denominator of D estimate to avoid issues during warm-up stage. True by default. "
+        "Ignored if optimizer is adamW",
+    )
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--report_to",
+        type=str,
+        default="tensorboard",
+        help=(
+            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
+            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    parser.add_argument(
+        "--upcast_before_saving",
+        action="store_true",
+        default=False,
+        help=(
+            "Whether to upcast the trained transformer layers to float32 before saving (at the end of training). "
+            "Defaults to precision dtype used for training to save memory"
+        ),
+    )
+    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
+
+    if input_args is not None:
+        args = parser.parse_args(input_args)
+    else:
+        args = parser.parse_args()
+
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    return args

utils/utils.py

@@ -0,0 +1,166 @@
+import os
+import shutil
+import threading
+from pathlib import Path
+
+import torch
+
+
+def import_from_transformers_modules(
+    pretrained_model_name_or_path, file_name, class_name
+):
+    import transformers
+
+    module_path = transformers.dynamic_module_utils.get_cached_module_file(
+        pretrained_model_name_or_path, file_name
+    )
+    return transformers.dynamic_module_utils.get_class_in_module(
+        class_name, module_path
+    )
+
+
+def deepspeed_zero_init_disabled_context_manager():
+    """
+    returns either a context list that includes one that will disable zero.Init or an empty context list
+    """
+    import accelerate
+
+    deepspeed_plugin = (
+        accelerate.state.AcceleratorState().deepspeed_plugin
+        if accelerate.state.is_initialized()
+        else None
+    )
+    if deepspeed_plugin is None:
+        return []
+
+    return [deepspeed_plugin.zero3_init_context_manager(enable=False)]
+
+
+def remove_excess_checkpoints(
+    save_directory,
+    checkpoints_total_limit: int = None,
+    checkpoint_prefix="checkpoint",
+    is_main_process: bool = True,
+):
+    # _after_ saving state, check if this save would set us over the `checkpoints_total_limit`
+    if is_main_process and checkpoints_total_limit is not None:
+        checkpoints = os.listdir(save_directory)
+        checkpoints = [d for d in checkpoints if d.startswith(checkpoint_prefix)]
+        checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[2]))
+
+        # _after_ we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit` checkpoints
+        if len(checkpoints) > checkpoints_total_limit:
+            num_to_remove = len(checkpoints) - checkpoints_total_limit
+            removing_checkpoints = checkpoints[0:num_to_remove]
+
+            print(
+                f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+            )
+            print(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+            for removing_checkpoint in removing_checkpoints:
+                removing_checkpoint = os.path.join(save_directory, removing_checkpoint)
+                shutil.rmtree(removing_checkpoint)
+
+
+def is_distributed_training():
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return True
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    return world_size > 1
+
+
+def contain_invalid_grad(optimizer):
+    invalid_grad = False
+    for param_group in optimizer.param_groups:
+        for param in param_group["params"]:
+            if param.grad is not None:
+                invalid_grad = invalid_grad or (
+                    torch.isnan(param.grad).any()
+                    or torch.isinf(param.grad).any()
+                    or torch.isneginf(param.grad).any()
+                )
+    if is_distributed_training():
+        invalid_grad_flag = torch.tensor(
+            [1.0 if invalid_grad else 0.0],
+            dtype=torch.float32,
+            requires_grad=False,
+        ).cuda()
+        torch.distributed.all_reduce(
+            invalid_grad_flag, op=torch.distributed.ReduceOp.MAX
+        )
+        invalid_grad = invalid_grad_flag.item() > 0
+    return invalid_grad
+
+
+def patch_npu_record_stream():
+    torch.utils.rename_privateuse1_backend("npu")
+    record_stream = torch.Tensor.record_stream
+
+    def _func(*args, **kwargs):
+        ret = record_stream(*args, **kwargs)
+        torch.cuda.synchronize()
+        return ret
+
+    torch.Tensor.record_stream = _func
+
+
+def patch_npu_diffusers_get_1d_rotary_pos_embed():
+    from typing import Union
+    import numpy as np
+    import diffusers
+
+    def __get_1d_rotary_pos_embed(
+        dim: int,
+        pos: Union[np.ndarray, int],
+        theta: float = 10000.0,
+        use_real=False,
+        linear_factor=1.0,
+        ntk_factor=1.0,
+        repeat_interleave_real=True,
+        freqs_dtype=torch.float32,  #  torch.float32, torch.float64 (flux)
+    ):
+        assert dim % 2 == 0
+
+        if isinstance(pos, int):
+            pos = torch.arange(pos)
+        if isinstance(pos, np.ndarray):
+            pos = torch.from_numpy(pos)  # type: ignore  # [S]
+
+        theta = theta * ntk_factor
+        freqs = (
+            1.0
+            / (
+                theta
+                ** (
+                    torch.arange(0, dim, 2, dtype=freqs_dtype, device=pos.device)[
+                        : (dim // 2)
+                    ]
+                    / dim
+                )
+            )
+            / linear_factor
+        )  # [D/2]
+        freqs = torch.outer(pos, freqs)  # type: ignore   # [S, D/2]
+        if use_real and repeat_interleave_real:
+            # flux, hunyuan-dit, cogvideox
+            freqs_cos = (
+                freqs.cos().float().repeat_interleave(2, dim=1).float()
+            )  # [S, D]
+            freqs_sin = (
+                freqs.sin().float().repeat_interleave(2, dim=1).float()
+            )  # [S, D]
+            return freqs_cos, freqs_sin
+        elif use_real:
+            # stable audio
+            freqs_cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1).float()  # [S, D]
+            freqs_sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1).float()  # [S, D]
+            return freqs_cos, freqs_sin
+        else:
+            # lumina
+            freqs_cis = torch.polar(
+                torch.ones_like(freqs), freqs
+            )  # complex64     # [S, D/2]
+            return freqs_cis
+
+    diffusers.models.embeddings.get_1d_rotary_pos_embed = __get_1d_rotary_pos_embed

utils/vprocess.py

@@ -0,0 +1,568 @@
+import base64
+import copy
+import logging
+import math
+import os
+import sys
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+from typing import Optional, Union, Tuple, List, Any, Dict
+from concurrent.futures import ThreadPoolExecutor
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+import numpy as np
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+PREFERRED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+
+def resizeinput(img):
+    multiple_of = 16
+    image_height, image_width = img.height, img.width
+    aspect_ratio = image_width / image_height
+    _, image_width, image_height = min(
+        (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
+    )
+    image_width = image_width // multiple_of * multiple_of
+    image_height = image_height // multiple_of * multiple_of
+    img = img.resize((image_width, image_height), Image.LANCZOS)
+    return img
+
+
+MAX_RATIO = 200
+SPATIAL_MERGE_SIZE = 2
+IMAGE_MIN_TOKEN_NUM = 4
+IMAGE_MAX_TOKEN_NUM = 16384
+VIDEO_MIN_TOKEN_NUM = 128
+VIDEO_MAX_TOKEN_NUM = 768
+
+FPS = 2.0
+FRAME_FACTOR = 2
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+MAX_NUM_WORKERS_FETCH_VIDEO = 8
+
+MODEL_SEQ_LEN = int(float(os.environ.get('MODEL_SEQ_LEN', 128000)))
+logger = logging.getLogger(__name__)
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(height: int, width: int, factor: int, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None) -> Tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    max_pixels = max_pixels if max_pixels is not None else (IMAGE_MAX_TOKEN_NUM * factor ** 2)
+    min_pixels = min_pixels if min_pixels is not None else (IMAGE_MIN_TOKEN_NUM * factor ** 2)
+    assert max_pixels >= min_pixels, "The max_pixels of image must be greater than or equal to min_pixels."
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def to_rgb(pil_image: Image.Image) -> Image.Image:
+      if pil_image.mode == 'RGBA':
+          white_background = Image.new("RGB", pil_image.size, (255, 255, 255))
+          white_background.paste(pil_image, mask=pil_image.split()[3])  # Use alpha channel as mask
+          return white_background
+      else:
+          return pil_image.convert("RGB")
+
+
+def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = 14) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+
+    image_obj = None
+    patch_factor = int(image_patch_size * SPATIAL_MERGE_SIZE)
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        with requests.get(image, stream=True) as response:
+            response.raise_for_status()
+            with BytesIO(response.content) as bio:
+                image_obj = copy.deepcopy(Image.open(bio))
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            with BytesIO(data) as bio:
+                image_obj = copy.deepcopy(Image.open(bio))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = to_rgb(image_obj)
+
+    ## resize
+    image = resizeinput(image)
+    # if "resized_height" in ele and "resized_width" in ele:
+    #     resized_height, resized_width = smart_resize(
+    #         ele["resized_height"],
+    #         ele["resized_width"],
+    #         factor=patch_factor,
+    #     )
+    # else:
+    #     width, height = image.size
+    #     min_pixels = ele.get("min_pixels", IMAGE_MIN_TOKEN_NUM * patch_factor ** 2)
+    #     max_pixels = ele.get("max_pixels", IMAGE_MAX_TOKEN_NUM * patch_factor ** 2)
+    #     resized_height, resized_width = smart_resize(
+    #         height,
+    #         width,
+    #         factor=patch_factor,
+    #         min_pixels=min_pixels,
+    #         max_pixels=max_pixels,
+    #     )
+    #     print(f"resized_height: {resized_height}, resized_width: {resized_width}")
+    # image = image.resize((resized_width, resized_height))
+    return image
+
+
+def smart_nframes(
+    ele: Dict[str, Any],
+    total_frames: int,
+    video_fps: Union[int, float],
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR)
+        nframes = total_frames / video_fps * fps
+        if nframes > total_frames:
+            logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]")
+        nframes = min(min(max(nframes, min_frames), max_frames), total_frames)
+        nframes = floor_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: Dict[str, Any],
+) -> Tuple[torch.Tensor, float]:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    video = video[idx]
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="torchvision",
+    )
+    return video, video_metadata, sample_fps
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def calculate_video_frame_range(
+    ele: Dict[str, Any],
+    total_frames: int,
+    video_fps: float,
+) -> Tuple[int, int, int]:
+    """
+    Calculate the start and end frame indices based on the given time range.
+
+    Args:
+        ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds).
+        total_frames (int): Total number of frames in the video.
+        video_fps (float): Frames per second of the video.
+
+    Returns:
+        tuple: A tuple containing (start_frame, end_frame, frame_count).
+
+    Raises:
+        ValueError: If input parameters are invalid or the time range is inconsistent.
+    """
+    # Validate essential parameters
+    if video_fps <= 0:
+        raise ValueError("video_fps must be a positive number")
+    if total_frames <= 0:
+        raise ValueError("total_frames must be a positive integer")
+
+    # Get start and end time in seconds
+    video_start = ele.get("video_start", None)
+    video_end = ele.get("video_end", None)
+    if video_start is None and video_end is None:
+        return 0, total_frames - 1, total_frames
+
+    max_duration = total_frames / video_fps
+    # Process start frame
+    if video_start is not None:
+        video_start_clamped = max(0.0, min(video_start, max_duration))
+        start_frame = math.ceil(video_start_clamped * video_fps)
+    else:
+        start_frame = 0
+    # Process end frame
+    if video_end is not None:
+        video_end_clamped = max(0.0, min(video_end, max_duration))
+        end_frame = math.floor(video_end_clamped * video_fps)
+        end_frame = min(end_frame, total_frames - 1)
+    else:
+        end_frame = total_frames - 1
+
+    # Validate frame order
+    if start_frame >= end_frame:
+        raise ValueError(
+            f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) "
+            f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). "
+            f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)"
+        )
+
+    logger.info(f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}")
+    return start_frame, end_frame, end_frame - start_frame + 1
+
+
+def _read_video_decord(
+    ele: Dict[str, Any],
+) -> Tuple[torch.Tensor, float]:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    start_frame, end_frame, total_frames = calculate_video_frame_range(
+        ele,
+        total_frames,
+        video_fps,
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="decord",
+    )
+    return video, video_metadata, sample_fps
+
+
+def is_torchcodec_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("torchcodec") is not None
+
+
+def _read_video_torchcodec(
+    ele: Dict[str, Any],
+) -> Tuple[torch.Tensor, float]:
+    """read video using torchcodec.decoders.VideoDecoder
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://", "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    from torchcodec.decoders import VideoDecoder
+    TORCHCODEC_NUM_THREADS = int(os.environ.get('TORCHCODEC_NUM_THREADS', 8))
+    logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}")
+    video_path = ele["video"]
+    st = time.time()
+    decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS)
+    video_fps = decoder.metadata.average_fps
+    total_frames = decoder.metadata.num_frames
+    start_frame, end_frame, total_frames = calculate_video_frame_range(
+        ele,
+        total_frames,
+        video_fps,
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist()
+    sample_fps = nframes / max(total_frames, 1e-6) * video_fps
+    video = decoder.get_frames_at(indices=idx).data
+    logger.info(f"torchcodec:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+
+    video_metadata = dict(
+        fps=video_fps,
+        frames_indices=idx,
+        total_num_frames=total_frames,
+        video_backend="torchcodec",
+    )
+    return video, video_metadata, sample_fps
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+    "torchcodec": _read_video_torchcodec,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_torchcodec_available():
+        video_reader_backend = "torchcodec"
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(ele: Dict[str, Any], image_patch_size: int = 14, return_video_sample_fps: bool = False,
+                return_video_metadata: bool = False) -> Union[torch.Tensor, List[Image.Image]]:
+    image_factor = image_patch_size * SPATIAL_MERGE_SIZE
+    VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor
+    VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        try:
+            video, video_metadata, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        except Exception as e:
+            logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}")
+            video, video_metadata, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele)
+    else:
+        # The input is a list of frames
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        # use ThreadPoolExecutor to parallel process frames
+        max_workers = min(MAX_NUM_WORKERS_FETCH_VIDEO, len(ele["video"]))
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            futures = [
+                executor.submit(fetch_image, {"image": video_element, **process_info}, image_factor)
+                for video_element in ele["video"]
+            ]
+            image_list = [future.result() for future in futures]
+
+        nframes = ceil_by_factor(len(image_list), FRAME_FACTOR)
+        if len(image_list) < nframes:
+            image_list.extend([image_list[-1]] * (nframes - len(image_list)))
+
+        sample_fps = ele.get("sample_fps", 2.0)
+        video = torch.stack([
+            torch.from_numpy(np.array(image).transpose(2, 0, 1))
+            for image in image_list
+        ])
+
+        # fake video metadata
+        raw_fps = process_info.pop("raw_fps", sample_fps)
+        video_metadata = dict(
+            fps=raw_fps,
+            frames_indices=[i for i in range(len(video))],
+            total_num_frames=(nframes / sample_fps) * raw_fps,
+        )
+
+    nframes, _, height, width = video.shape
+    min_pixels = ele.get("min_pixels", VIDEO_FRAME_MIN_PIXELS)
+    total_pixels = ele.get("total_pixels", MODEL_SEQ_LEN * image_factor * image_factor * 0.9)
+    max_pixels = max(min(VIDEO_FRAME_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05))
+    max_pixels_supposed = ele.get("max_pixels", max_pixels)
+    if max_pixels_supposed > max_pixels:
+        logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].")
+    max_pixels = min(max_pixels_supposed, max_pixels)
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=image_factor,
+        )
+    else:
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=image_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    video = transforms.functional.resize(
+        video,
+        [resized_height, resized_width],
+        interpolation=InterpolationMode.BICUBIC,
+        antialias=True,
+    ).float()
+
+    final_video = (video, video_metadata) if return_video_metadata else video
+    if return_video_sample_fps:
+        return final_video, sample_fps
+    return final_video
+
+
+def extract_vision_info(conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]]) -> List[Dict[str, Any]]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele.get("type", "text") in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]],
+    return_video_kwargs: bool = False,
+    return_video_metadata: bool = False,
+    image_patch_size: int = 14,
+) -> Tuple[Optional[List[Image.Image]], Optional[List[Union[torch.Tensor, List[Image.Image]]]], Optional[Dict[str, Any]]]:
+
+    vision_infos = extract_vision_info(conversations)
+    ## Read images or videos
+    image_inputs = []
+    video_inputs = []
+    video_sample_fps_list = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info, image_patch_size=image_patch_size))
+        elif "video" in vision_info:
+            video_input, video_sample_fps = fetch_video(vision_info, return_video_sample_fps=True,
+                        image_patch_size=image_patch_size, return_video_metadata=return_video_metadata)
+            video_sample_fps_list.append(video_sample_fps)
+            video_inputs.append(video_input)
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+
+    video_kwargs = {'do_sample_frames': False}
+    if not return_video_metadata: # BC for qwen2.5vl
+        video_kwargs.update({'fps': video_sample_fps_list})
+
+    if return_video_kwargs:
+        return image_inputs, video_inputs, video_kwargs
+    return image_inputs, video_inputs

web_edit.py

@@ -0,0 +1,252 @@
+import torch
+from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline
+from diffusers.utils import load_image
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+# from qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
+import os
+import re
+from PIL import Image
+import gradio as gr
+import uuid
+import argparse 
+
+def parse_args():
+    """Parses command-line arguments for model paths and server configuration."""
+    parser = argparse.ArgumentParser(description="Launch DreamOmni2 Editing Gradio Demo.")
+    parser.add_argument(
+        "--vlm_path", 
+        type=str, 
+        default="vlm-model", 
+        help="Path to the Qwen2_5_VL VLM model directory."
+    )
+    parser.add_argument(
+        "--edit_lora_path", 
+        type=str, 
+        default="edit_lora", 
+        help="Path to the FLUX.1-Kontext editing LoRA weights directory."
+    )
+    parser.add_argument(
+        "--server_name", 
+        type=str, 
+        default="0.0.0.0", 
+        help="The server name (IP address) to host the Gradio demo."
+    )
+    parser.add_argument(
+        "--server_port", 
+        type=int, 
+        default=7860, 
+        help="The port number to host the Gradio demo."
+    )
+    args = parser.parse_args()
+    return args
+
+ARGS = parse_args()
+vlm_path = ARGS.vlm_path
+edit_lora_path = ARGS.edit_lora_path
+server_name = ARGS.server_name
+server_port = ARGS.server_port
+device = "cuda"
+
+def extract_gen_content(text):
+    text = text[6:-7]
+    return text
+
+print(f"Loading models from vlm_path: {vlm_path}, edit_lora_path: {edit_lora_path}")
+
+pipe = DreamOmni2Pipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Kontext-dev",
+    torch_dtype=torch.bfloat16
+)
+pipe.to(device)
+pipe.load_lora_weights(edit_lora_path, adapter_name="edit")
+pipe.set_adapters(["edit"], adapter_weights=[1])    
+
+vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    vlm_path,
+    torch_dtype="bfloat16",
+    device_map="cuda"
+)
+processor = AutoProcessor.from_pretrained(vlm_path)
+
+
+def infer_vlm(input_img_path, input_instruction, prefix):
+    if not vlm_model or not processor:
+        raise gr.Error("VLM Model not loaded. Cannot process prompt.")
+    tp = []
+    for path in input_img_path:
+        tp.append({"type": "image", "image": path})
+    tp.append({"type": "text", "text": input_instruction + prefix})
+    messages = [{"role": "user", "content": tp}]
+
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
+    inputs = inputs.to("cuda")
+
+    generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    return output_text[0]
+
+PREFERRED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+def find_closest_resolution(width, height, preferred_resolutions):
+    input_ratio = width / height
+    closest_resolution = min(
+        preferred_resolutions,
+        key=lambda res: abs((res[0] / res[1]) - input_ratio)
+    )
+    return closest_resolution
+
+def perform_edit(input_img_paths, input_instruction, output_path):
+    prefix = " It is editing task."
+    source_imgs = []
+    for path in input_img_paths:
+        img = load_image(path)
+        # source_imgs.append(img)
+        source_imgs.append(resizeinput(img))
+    prompt = infer_vlm(input_img_paths, input_instruction, prefix)
+    prompt = extract_gen_content(prompt)
+    print(f"Generated Prompt for VLM: {prompt}")
+
+    image = pipe(
+        images=source_imgs,
+        height=source_imgs[0].height,
+        width=source_imgs[0].width,
+        prompt=prompt,
+        num_inference_steps=30,
+        guidance_scale=3.5,
+    ).images[0]
+    image.save(output_path)
+    print(f"Edit result saved to {output_path}")
+
+
+def process_request(image_file_1, image_file_2, instruction):
+    # debugpy.listen(5678)
+    # print("Waiting for debugger attach...")
+    # debugpy.wait_for_client()
+    if not image_file_1 or not image_file_2:
+        raise gr.Error("Please upload both images.")
+    if not instruction:
+        raise gr.Error("Please provide an instruction.")
+    if not pipe or not vlm_model:
+        raise gr.Error("Models not loaded. Check the console for errors.")
+    
+    output_path = f"/tmp/{uuid.uuid4()}.png"
+    input_img_paths = [image_file_1, image_file_2]  # List of file paths from the two gr.File inputs
+
+    perform_edit(input_img_paths, instruction, output_path)
+    return output_path
+
+
+css = """
+.text-center { text-align: center; }
+.result-img img {
+    max-height: 60vh !important; 
+    min-height: 30vh !important;
+    width: auto !important;      
+    object-fit: contain;         
+}
+.input-img img {
+    max-height: 30vh !important; 
+    width: auto !important;      
+    object-fit: contain;         
+}
+"""
+
+
+with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo:
+    gr.HTML(
+        """
+        <h1 style="text-align:center; font-size:48px; font-weight:bold; margin-bottom:20px;">
+            DreamOmni2: Omni-purpose Image Generation and Editing
+        </h1>
+        """
+    )
+    gr.Markdown(
+        "Select a mode, upload two images, provide an instruction, and click 'Run'.",
+        elem_classes="text-center"
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("⬆️ Upload images. Click or drag to upload.")
+            
+            with gr.Row():
+                image_uploader_1 = gr.Image(
+                    label="Img 1",
+                    type="filepath",
+                    interactive=True,
+                    elem_classes="input-img",
+                )
+                image_uploader_2 = gr.Image(
+                    label="Img 2",
+                    type="filepath",
+                    interactive=True,
+                    elem_classes="input-img",
+                )
+            
+            instruction_text = gr.Textbox(
+                label="Instruction",
+                lines=2,
+                placeholder="Input your instruction for generation or editing here...",
+            )
+            run_button = gr.Button("Run", variant="primary")
+
+        with gr.Column(scale=2):
+            gr.Markdown(
+                "✏️ **Editing Mode**: Modify an existing image using instructions and references.\n\n"
+                "Tip: If the result is not what you expect, try clicking **Run** again. "
+            )
+            output_image = gr.Image(
+                label="Result",
+                type="filepath",
+                elem_classes="result-img",
+            )
+
+    # --- Examples (不变) ---
+    gr.Markdown("## Examples")
+
+    gr.Examples(
+        label="Editing Examples",
+        examples=[
+            ["example_input/edit_tests/4/ref_0.jpg", "example_input/edit_tests/4/ref_1.jpg", "Replace the first image have the same image style as the second image.","example_input/edit_tests/4/res.jpg"],
+            ["example_input/edit_tests/5/ref_0.jpg", "example_input/edit_tests/5/ref_1.jpg", "Make the person in the first image have the same hairstyle as the person in the second image.","example_input/edit_tests/5/res.jpg"],
+            ["example_input/edit_tests/src.jpg", "example_input/edit_tests/ref.jpg", "Make the woman from the second image stand on the road in the first image.","example_input/edit_tests/edi_res.png"],
+            ["example_input/edit_tests/1/ref_0.jpg", "example_input/edit_tests/1/ref_1.jpg", "Replace the lantern in the first image with the dog in the second image.","example_input/edit_tests/1/res.jpg"],
+            ["example_input/edit_tests/2/ref_0.jpg", "example_input/edit_tests/2/ref_1.jpg", "Replace the suit in the first image with the clothes in the second image.","example_input/edit_tests/2/res.jpg"],
+            ["example_input/edit_tests/3/ref_0.jpg", "example_input/edit_tests/3/ref_1.jpg", "Make the first image has the same light condition as the second image.","example_input/edit_tests/3/res.jpg"],
+            ["example_input/edit_tests/6/ref_0.jpg", "example_input/edit_tests/6/ref_1.jpg", "Make the words in the first image have the same font as the words in the second image.","example_input/edit_tests/6/res.jpg"],
+            ["example_input/edit_tests/7/ref_0.jpg", "example_input/edit_tests/7/ref_1.jpg", "Make the car in the first image have the same pattern as the mouse in the second image.","example_input/edit_tests/7/res.jpg"],
+            ["example_input/edit_tests/8/ref_0.jpg", "example_input/edit_tests/8/ref_1.jpg", "Make the dress in the first image have the same pattern in the second image.","example_input/edit_tests/8/res.jpg"],
+        ],
+        inputs=[image_uploader_1, image_uploader_2, instruction_text, output_image],
+        cache_examples=False,
+    )
+
+    run_button.click(
+        fn=process_request,
+        inputs=[image_uploader_1, image_uploader_2, instruction_text],
+        outputs=output_image
+    )
+
+if __name__ == "__main__":
+    print("Launching Gradio Demo...")
+    demo.launch(server_name=server_name, server_port=server_port)

web_generate.py

@@ -0,0 +1,251 @@
+import torch
+from dreamomni2.pipeline_dreamomni2 import DreamOmni2Pipeline
+from diffusers.utils import load_image
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+# from qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
+import os
+import re
+from PIL import Image
+import gradio as gr
+import uuid
+import argparse 
+
+def parse_args():
+    """Parses command-line arguments for model paths and server configuration."""
+    parser = argparse.ArgumentParser(description="Launch DreamOmni2 Editing Gradio Demo.")
+    parser.add_argument(
+        "--vlm_path", 
+        type=str, 
+        default="vlm-model", 
+        help="Path to the Qwen2_5_VL VLM model directory."
+    )
+    parser.add_argument(
+        "--gen_lora_path", 
+        type=str, 
+        default="gen_lora", 
+        help="Path to the FLUX.1-Kontext generation LoRA weights directory."
+    )
+    parser.add_argument(
+        "--server_name", 
+        type=str, 
+        default="0.0.0.0", 
+        help="The server name (IP address) to host the Gradio demo."
+    )
+    parser.add_argument(
+        "--server_port", 
+        type=int, 
+        default=7860, 
+        help="The port number to host the Gradio demo."
+    )
+    args = parser.parse_args()
+    return args
+
+ARGS = parse_args()
+vlm_path = ARGS.vlm_path
+gen_lora_path = ARGS.gen_lora_path
+server_name = ARGS.server_name
+server_port = ARGS.server_port
+device = "cuda"
+
+def extract_gen_content(text):
+    text = text[6:-7]
+    return text
+
+print(f"Loading models from vlm_path: {vlm_path}, gen_lora_path: {gen_lora_path}")
+
+pipe = DreamOmni2Pipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Kontext-dev",
+    torch_dtype=torch.bfloat16
+)
+pipe.to(device)
+pipe.load_lora_weights(gen_lora_path, adapter_name="generation")
+pipe.set_adapters(["generation"], adapter_weights=[1])  
+
+vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    vlm_path,
+    torch_dtype="bfloat16",
+    device_map="cuda"
+)
+processor = AutoProcessor.from_pretrained(vlm_path)
+
+
+def infer_vlm(input_img_path, input_instruction, prefix):
+    if not vlm_model or not processor:
+        raise gr.Error("VLM Model not loaded. Cannot process prompt.")
+    tp = []
+    for path in input_img_path:
+        tp.append({"type": "image", "image": path})
+    tp.append({"type": "text", "text": input_instruction + prefix})
+    messages = [{"role": "user", "content": tp}]
+
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
+    inputs = inputs.to("cuda")
+
+    generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    return output_text[0]
+
+
+PREFERRED_KONTEXT_RESOLUTIONS = [
+    (672, 1568),
+    (688, 1504),
+    (720, 1456),
+    (752, 1392),
+    (800, 1328),
+    (832, 1248),
+    (880, 1184),
+    (944, 1104),
+    (1024, 1024),
+    (1104, 944),
+    (1184, 880),
+    (1248, 832),
+    (1328, 800),
+    (1392, 752),
+    (1456, 720),
+    (1504, 688),
+    (1568, 672),
+]
+def find_closest_resolution(width, height, preferred_resolutions):
+    input_ratio = width / height
+    closest_resolution = min(
+        preferred_resolutions,
+        key=lambda res: abs((res[0] / res[1]) - input_ratio)
+    )
+    return closest_resolution
+
+def perform_generation(input_img_paths, input_instruction, output_path, height=1024, width=1024):
+    prefix = " It is generation task."
+    source_imgs = []
+    for path in input_img_paths:
+        img = load_image(path)
+        # source_imgs.append(img)
+        source_imgs.append(resizeinput(img))
+    prompt = infer_vlm(input_img_paths, input_instruction, prefix)
+    prompt = extract_gen_content(prompt)
+    print(f"Generated Prompt for VLM: {prompt}")
+    
+    image = pipe(
+        images=source_imgs,
+        height=height,
+        width=width,
+        prompt=prompt,
+        num_inference_steps=30,
+        guidance_scale=3.5,
+    ).images[0]
+    
+    image.save(output_path)
+    print(f"Generation result saved to {output_path}")
+
+# --- Gradio Interface Logic ---
+
+def process_request(image_file_1, image_file_2, instruction):
+    # debugpy.listen(5678)
+    # print("Waiting for debugger attach...")
+    # debugpy.wait_for_client()
+    if not image_file_1 or not image_file_2:
+        raise gr.Error("Please upload both images.")
+    if not instruction:
+        raise gr.Error("Please provide an instruction.")
+    if not pipe or not vlm_model:
+        raise gr.Error("Models not loaded. Check the console for errors.")
+    
+    output_path = f"/tmp/{uuid.uuid4()}.png"
+    input_img_paths = [image_file_1, image_file_2]  # List of file paths from the two gr.File inputs
+
+    perform_generation(input_img_paths, instruction, output_path)
+    
+    return output_path
+
+
+css = """
+.text-center { text-align: center; }
+.result-img img {
+    max-height: 60vh !important; 
+    min-height: 30vh !important;
+    width: auto !important;      
+    object-fit: contain;         
+}
+.input-img img {
+    max-height: 30vh !important; 
+    width: auto !important;      
+    object-fit: contain;         
+}
+"""
+
+
+with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo:
+    gr.HTML(
+        """
+        <h1 style="text-align:center; font-size:48px; font-weight:bold; margin-bottom:20px;">
+            DreamOmni2: Omni-purpose Image Generation and Editing
+        </h1>
+        """
+    )
+    gr.Markdown(
+        "Select a mode, upload two images, provide an instruction, and click 'Run'.",
+        elem_classes="text-center"
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.Markdown("⬆️ Upload images. Click or drag to upload.")
+            
+            with gr.Row():
+                image_uploader_1 = gr.Image(
+                    label="Img 1",
+                    type="filepath",
+                    interactive=True,
+                    elem_classes="input-img",
+                )
+                image_uploader_2 = gr.Image(
+                    label="Img 2",
+                    type="filepath",
+                    interactive=True,
+                    elem_classes="input-img",
+                )
+            
+            instruction_text = gr.Textbox(
+                label="Instruction",
+                lines=2,
+                placeholder="Input your instruction for generation or editing here...",
+            )
+            run_button = gr.Button("Run", variant="primary")
+
+        with gr.Column(scale=2):
+            gr.Markdown("🖼️ **Generation Mode**: Create new scenes from reference images."
+                        "Tip: If the result is not what you expect, try clicking **Run** again. ")
+            output_image = gr.Image(
+                label="Result",
+                type="filepath",
+                elem_classes="result-img",
+            )
+
+    # --- Examples ---
+    gr.Markdown("## Examples")
+    gr.Examples(
+        label="Generation Examples",
+        examples=[
+            [
+                "example_input/gen_tests/img1.jpg",
+                "example_input/gen_tests/img2.jpg",
+                "In the scene, the character from the first image stands on the left, and the character from the second image stands on the right. They are shaking hands against the backdrop of a spaceship interior.",
+                "example_input/gen_tests/gen_res.png"
+            ]
+        ],
+        inputs=[image_uploader_1, image_uploader_2, instruction_text, output_image],
+        cache_examples=False,
+    )
+
+    run_button.click(
+        fn=process_request,
+        inputs=[image_uploader_1, image_uploader_2, instruction_text],
+        outputs=output_image
+    )
+
+if __name__ == "__main__":
+    
+    print("Launching Gradio Demo...")
+    demo.launch(server_name="0.0.0.0", server_port=7861, )