update code

.gitignore

@@ -1,220 +0,0 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[codz]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#   Usually these files are written by a python script from a template
-#   before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py.cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-# Pipfile.lock
-
-# UV
-#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-# uv.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-# poetry.lock
-# poetry.toml
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
-#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
-# pdm.lock
-# pdm.toml
-.pdm-python
-.pdm-build/
-
-# pixi
-#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
-# pixi.lock
-#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
-#   in the .venv directory. It is recommended not to include this directory in version control.
-.pixi
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# Redis
-*.rdb
-*.aof
-*.pid
-
-# RabbitMQ
-mnesia/
-rabbitmq/
-rabbitmq-data/
-
-# ActiveMQ
-activemq-data/
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.envrc
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#   and can be added to the global gitignore or merged into this file.  For a more nuclear
-#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
-# .idea/
-
-# Abstra
-#   Abstra is an AI-powered process automation framework.
-#   Ignore directories containing user credentials, local state, and settings.
-#   Learn more at https://abstra.io/docs
-.abstra/
-
-# Visual Studio Code
-#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
-#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
-#   you could uncomment the following to ignore the entire vscode folder
-# .vscode/
-
-# Ruff stuff:
-.ruff_cache/
-
-# PyPI configuration file
-.pypirc
-
-# Marimo
-marimo/_static/
-marimo/_lsp/
-__marimo__/
-
-# Streamlit
-.streamlit/secrets.toml
-
-tmp_script.sh
-.gradio/
-gradio_tmp

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.

app.py

@@ -8,12 +8,16 @@ import argparse
 from huggingface_hub import login, snapshot_download
 
 import torch
-from pipeline_flux_kontext import FluxKontextPipeline
+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
 
 
+def extract_gen_content(text):
+    text = text[6:-7]
+    return text
+
 def _load_model_processor():
 
     local_dir = snapshot_download(
@@ -21,22 +25,22 @@ def _load_model_processor():
         revision="main",
         allow_patterns=["vlm-model/**", "edit_lora/**"],
     )
-    local_vlm_dir = os.path.join(local_dir, 'vlm-model')
-    local_lora_dir = os.path.join(local_dir, 'edit_lora')
+    vlm_dir = os.path.join(local_dir, 'vlm-model')
+    lora_dir = os.path.join(local_dir, 'edit_lora')
 
-    print(f"Loading models from vlm_path: {local_vlm_dir}, edit_lora_path: {local_lora_dir}")
-    pipe = FluxKontextPipeline.from_pretrained(
+    print(f"Loading models from vlm_path: {vlm_dir}, edit_lora_path: {lora_dir}")
+    pipe = DreamOmni2Pipeline.from_pretrained(
         "black-forest-labs/FLUX.1-Kontext-dev",
         torch_dtype=torch.bfloat16
     )
-    pipe.load_lora_weights(local_lora_dir, adapter_name="edit")
+    pipe.load_lora_weights(lora_dir, adapter_name="edit")
     pipe.set_adapters(["edit"], adapter_weights=[1])    
 
     vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-        local_vlm_dir,
+        vlm_dir,
         torch_dtype="bfloat16"
     )
-    processor = AutoProcessor.from_pretrained(local_vlm_dir)
+    processor = AutoProcessor.from_pretrained(vlm_dir)
     return vlm_model, processor, pipe
 
 
@@ -88,28 +92,23 @@ def _launch_demo(vlm_model, processor, pipe):
             key=lambda res: abs((res[0] / res[1]) - input_ratio)
         )
         return closest_resolution
-    
-    def extract_gen_content(text):
-        text = text[6:-7]
-        return text
 
     @spaces.GPU()
     def perform_edit(input_img_paths, input_instruction, output_path):
         prefix = " It is editing task."
-        source_imgs = [load_image(path) for path in input_img_paths]
-        resized_imgs = []
-        for img in source_imgs:
-            target_resolution = find_closest_resolution(img.width, img.height, PREFERRED_KONTEXT_RESOLUTIONS)
-            resized_img = img.resize(target_resolution, Image.LANCZOS)
-            resized_imgs.append(resized_img)
+        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=resized_imgs,
-            height=resized_imgs[0].height,
-            width=resized_imgs[0].width,
+            images=source_imgs,
+            height=source_imgs[0].height,
+            width=source_imgs[0].width,
             prompt=prompt,
             num_inference_steps=30,
             guidance_scale=3.5,
@@ -205,15 +204,15 @@ def _launch_demo(vlm_model, processor, pipe):
         gr.Examples(
             label="Editing Examples",
             examples=[
-                ["edit_tests/4/ref_0.jpg", "edit_tests/4/ref_1.jpg", "Replace the first image have the same image style as the second image.","edit_tests/4/res.jpg"],
-                ["edit_tests/5/ref_0.jpg", "edit_tests/5/ref_1.jpg", "Make the person in the first image have the same hairstyle as the person in the second image.","edit_tests/5/res.jpg"],
-                ["edit_tests/src.jpg", "edit_tests/ref.jpg", "Make the woman from the second image stand on the road in the first image.","edit_tests/edi_res.png"],
-                ["edit_tests/1/ref_0.jpg", "edit_tests/1/ref_1.jpg", "Replace the lantern in the first image with the dog in the second image.","edit_tests/1/res.jpg"],
-                ["edit_tests/2/ref_0.jpg", "edit_tests/2/ref_1.jpg", "Replace the suit in the first image with the clothes in the second image.","edit_tests/2/res.jpg"],
-                ["edit_tests/3/ref_0.jpg", "edit_tests/3/ref_1.jpg", "Make the first image has the same light condition as the second image.","edit_tests/3/res.jpg"],
-                ["edit_tests/6/ref_0.jpg", "edit_tests/6/ref_1.jpg", "Make the words in the first image have the same font as the words in the second image.","edit_tests/6/res.jpg"],
-                ["edit_tests/7/ref_0.jpg", "edit_tests/7/ref_1.jpg", "Make the car in the first image have the same pattern as the mouse in the second image.","edit_tests/7/res.jpg"],
-                ["edit_tests/8/ref_0.jpg", "edit_tests/8/ref_1.jpg", "Make the dress in the first image have the same pattern in the second image.","edit_tests/8/res.jpg"],
+                ["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,

pipeline_flux_kontext.py → dreamomni2/pipeline_dreamomni2.py

@@ -57,10 +57,10 @@ EXAMPLE_DOC_STRING = """
     Examples:
         ```py
         >>> import torch
-        >>> from diffusers import FluxKontextPipeline
+        >>> from diffusers import DreamOmni2Pipeline
         >>> from diffusers.utils import load_image
 
-        >>> pipe = FluxKontextPipeline.from_pretrained(
+        >>> pipe = DreamOmni2Pipeline.from_pretrained(
         ...     "black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16
         ... )
         >>> pipe.to("cuda")
@@ -187,7 +187,7 @@ def retrieve_latents(
         raise AttributeError("Could not access latents of provided encoder_output")
 
 
-class FluxKontextPipeline(
+class DreamOmni2Pipeline(
     DiffusionPipeline,
     FluxLoraLoaderMixin,
     FromSingleFileMixin,
@@ -694,10 +694,7 @@ class FluxKontextPipeline(
                 image_ids = self._prepare_latent_image_ids(
                     batch_size, image_latent_height // 2, image_latent_width // 2, device, dtype
                 )
-                # image ids are the same as latent ids with the first dimension set to 1 instead of 0
-                # image_ids[..., 0] = 0.9+i*0.1 
                 image_ids[..., 0] = i+1
-                # image_ids[..., 1] += h_offset
                 image_ids[..., 2] += w_offset
                 tp_image_latents.append(image_latents)
                 tp_image_ids.append(image_ids)

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)

script.sh

@@ -1,12 +0,0 @@
-CUDA_VISIBLE_DEVICES=3 GRADIO_TEMP_DIR=gradio_tmp python web_edit.py \
-    --vlm_path /gpfs/bhpeng/generation/do2/vlm-model \
-    --edit_lora_path /gpfs/bhpeng/generation/do2/edit_lora \
-    --server_name "0.0.0.0" \
-    --server_port 7869
-
-
-CUDA_VISIBLE_DEVICES=1 python web_generate.py \
-    --vlm_path /gpfs/bhpeng/generation/do2/vlm-model \
-    --gen_lora_path /gpfs/bhpeng/generation/do2/gen_lora \
-    --server_name "0.0.0.0" \
-    --server_port 7861

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

@@ -1,8 +1,9 @@
 import torch
-from pipeline_flux_kontext import FluxKontextPipeline
+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 qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
 import os
 import re
 from PIL import Image
@@ -53,7 +54,7 @@ def extract_gen_content(text):
 
 print(f"Loading models from vlm_path: {vlm_path}, edit_lora_path: {edit_lora_path}")
 
-pipe = FluxKontextPipeline.from_pretrained(
+pipe = DreamOmni2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Kontext-dev",
     torch_dtype=torch.bfloat16
 )
@@ -117,20 +118,19 @@ def find_closest_resolution(width, height, preferred_resolutions):
 
 def perform_edit(input_img_paths, input_instruction, output_path):
     prefix = " It is editing task."
-    source_imgs = [load_image(path) for path in input_img_paths]
-    resized_imgs = []
-    for img in source_imgs:
-        target_resolution = find_closest_resolution(img.width, img.height, PREFERRED_KONTEXT_RESOLUTIONS)
-        resized_img = img.resize(target_resolution, Image.LANCZOS)
-        resized_imgs.append(resized_img)
+    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=resized_imgs,
-        height=resized_imgs[0].height,
-        width=resized_imgs[0].width,
+        images=source_imgs,
+        height=source_imgs[0].height,
+        width=source_imgs[0].width,
         prompt=prompt,
         num_inference_steps=30,
         guidance_scale=3.5,
@@ -227,15 +227,15 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo:
     gr.Examples(
         label="Editing Examples",
         examples=[
-            ["edit_tests/4/ref_0.jpg", "edit_tests/4/ref_1.jpg", "Replace the first image have the same image style as the second image.","edit_tests/4/res.jpg"],
-            ["edit_tests/5/ref_0.jpg", "edit_tests/5/ref_1.jpg", "Make the person in the first image have the same hairstyle as the person in the second image.","edit_tests/5/res.jpg"],
-            ["edit_tests/src.jpg", "edit_tests/ref.jpg", "Make the woman from the second image stand on the road in the first image.","edit_tests/edi_res.png"],
-            ["edit_tests/1/ref_0.jpg", "edit_tests/1/ref_1.jpg", "Replace the lantern in the first image with the dog in the second image.","edit_tests/1/res.jpg"],
-            ["edit_tests/2/ref_0.jpg", "edit_tests/2/ref_1.jpg", "Replace the suit in the first image with the clothes in the second image.","edit_tests/2/res.jpg"],
-            ["edit_tests/3/ref_0.jpg", "edit_tests/3/ref_1.jpg", "Make the first image has the same light condition as the second image.","edit_tests/3/res.jpg"],
-            ["edit_tests/6/ref_0.jpg", "edit_tests/6/ref_1.jpg", "Make the words in the first image have the same font as the words in the second image.","edit_tests/6/res.jpg"],
-            ["edit_tests/7/ref_0.jpg", "edit_tests/7/ref_1.jpg", "Make the car in the first image have the same pattern as the mouse in the second image.","edit_tests/7/res.jpg"],
-            ["edit_tests/8/ref_0.jpg", "edit_tests/8/ref_1.jpg", "Make the dress in the first image have the same pattern in the second image.","edit_tests/8/res.jpg"],
+            ["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,

web_generate.py

@@ -1,8 +1,9 @@
 import torch
-from pipeline_flux_kontext import FluxKontextPipeline
+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 qwen_vl_utils import process_vision_info
+from utils.vprocess import process_vision_info, resizeinput
 import os
 import re
 from PIL import Image
@@ -53,7 +54,7 @@ def extract_gen_content(text):
 
 print(f"Loading models from vlm_path: {vlm_path}, gen_lora_path: {gen_lora_path}")
 
-pipe = FluxKontextPipeline.from_pretrained(
+pipe = DreamOmni2Pipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Kontext-dev",
     torch_dtype=torch.bfloat16
 )
@@ -118,18 +119,17 @@ def find_closest_resolution(width, height, preferred_resolutions):
 
 def perform_generation(input_img_paths, input_instruction, output_path, height=1024, width=1024):
     prefix = " It is generation task."
-    source_imgs = [load_image(path) for path in input_img_paths]
-    resized_imgs = []
-    for img in source_imgs:
-        target_resolution = find_closest_resolution(img.width, img.height, PREFERRED_KONTEXT_RESOLUTIONS)
-        resized_img = img.resize(target_resolution, Image.LANCZOS)
-        resized_imgs.append(resized_img)
+    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=resized_imgs,
+        images=source_imgs,
         height=height,
         width=width,
         prompt=prompt,
@@ -229,10 +229,10 @@ with gr.Blocks(theme=gr.themes.Soft(), title="DreamOmni2", css=css) as demo:
         label="Generation Examples",
         examples=[
             [
-                "gen_tests/img1.jpg",
-                "gen_tests/img2.jpg",
+                "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.",
-                "gen_tests/gen_res.png"
+                "example_input/gen_tests/gen_res.png"
             ]
         ],
         inputs=[image_uploader_1, image_uploader_2, instruction_text, output_image],