diff --git a/.DS_Store b/.DS_Store
deleted file mode 100644
index 84f5661..0000000
Binary files a/.DS_Store and /dev/null differ
diff --git a/.gitignore b/.gitignore
index b7faf40..dfa1df5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -205,3 +205,30 @@ cython_debug/
marimo/_static/
marimo/_lsp/
__marimo__/
+
+
+
+pretrained_model
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+*.json
+*.jsonl
+env/
+venv/
+.venv/
+results/
+fusion_result.json
+*output*
+*kernel_meta/
+data/*.json
+maps
+data/metric_cache
+data/maps
+tmp
+navsim_v1
+nuscenes
+extra-info/
+pretrained_model/
+.DS_Store
\ No newline at end of file
diff --git a/.pylintrc b/.pylintrc
index 7e930ac..aa823bf 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -31,7 +31,7 @@ extension-pkg-allow-list=
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code. (This is an alternative name to extension-pkg-allow-list
# for backward compatibility.)
-extension-pkg-whitelist=cv2
+extension-pkg-whitelist=
# Return non-zero exit code if any of these messages/categories are detected,
# even if score is above --fail-under value. Syntax same as enable. Messages
@@ -59,15 +59,16 @@ ignore-paths=
# Emacs file locks
ignore-patterns=^\.#
-# List of module names for which member attributes should not be checked
-# (useful for modules/projects where namespaces are manipulated during runtime
-# and thus existing member attributes cannot be deduced by static analysis). It
-# supports qualified module names, as well as Unix pattern matching.
-ignored-modules=cv2
+# List of module names for which member attributes should not be checked and
+# will not be imported (useful for modules/projects where namespaces are
+# manipulated during runtime and thus existing member attributes cannot be
+# deduced by static analysis). It supports qualified module names, as well as
+# Unix pattern matching.
+ignored-modules=torch
# Python code to execute, usually for sys.path manipulation such as
# pygtk.require().
-init-hook='import sys; sys.path.append(".")'
+#init-hook=
# Use multiple processes to speed up Pylint. Specifying 0 will auto-detect the
# number of processors available to use, and will cap the count on Windows to
@@ -86,9 +87,13 @@ load-plugins=
# Pickle collected data for later comparisons.
persistent=yes
+# Resolve imports to .pyi stubs if available. May reduce no-member messages and
+# increase not-an-iterable messages.
+prefer-stubs=no
+
# Minimum Python version to use for version dependent checks. Will default to
# the version used to run pylint.
-py-version=3.10
+py-version=3.12
# Discover python modules and packages in the file system subtree.
recursive=no
@@ -99,10 +104,6 @@ recursive=no
# source root.
source-roots=
-# When enabled, pylint would attempt to guess common misconfiguration and emit
-# user-friendly hints instead of false-positive error messages.
-suggestion-mode=yes
-
# Allow loading of arbitrary C extensions. Extensions are imported into the
# active Python interpreter and may run arbitrary code.
unsafe-load-any-extension=no
@@ -229,6 +230,11 @@ name-group=
# not require a docstring.
no-docstring-rgx=^_
+# Regular expression matching correct parameter specification variable names.
+# If left empty, parameter specification variable names will be checked with
+# the set naming style.
+#paramspec-rgx=
+
# List of decorators that produce properties, such as abc.abstractproperty. Add
# to this list to register other decorators that produce valid properties.
# These decorators are taken in consideration only for invalid-name.
@@ -242,13 +248,17 @@ property-classes=abc.abstractproperty
# variable names will be checked with the set naming style.
#typevar-rgx=
+# Regular expression matching correct type variable tuple names. If left empty,
+# type variable tuple names will be checked with the set naming style.
+#typevartuple-rgx=
+
# Naming style matching correct variable names.
variable-naming-style=snake_case
# Regular expression matching correct variable names. Overrides variable-
# naming-style. If left empty, variable names will be checked with the set
# naming style.
-variable-rgx=(_?[a-z][A-Za-z0-9]{0,30})|([A-Z0-9]{1,30})
+#variable-rgx=
[CLASSES]
@@ -285,23 +295,26 @@ exclude-too-few-public-methods=
ignored-parents=
# Maximum number of arguments for function / method.
-max-args=7
+max-args=10
# Maximum number of attributes for a class (see R0902).
-max-attributes=20
+max-attributes=7
# Maximum number of boolean expressions in an if statement (see R0916).
max-bool-expr=5
# Maximum number of branch for function / method body.
-max-branches=12
+max-branches=40
# Maximum number of locals for function / method body.
-max-locals=15
+max-locals=50
# Maximum number of parents for a class (see R0901).
max-parents=7
+# Maximum number of positional arguments for function / method.
+max-positional-arguments=10
+
# Maximum number of public methods for a class (see R0904).
max-public-methods=20
@@ -309,10 +322,10 @@ max-public-methods=20
max-returns=6
# Maximum number of statements in function / method body.
-max-statements=300
+max-statements=150
# Minimum number of public methods for a class (see R0903).
-min-public-methods=1
+min-public-methods=2
[EXCEPTIONS]
@@ -336,11 +349,13 @@ indent-after-paren=4
# tab).
indent-string=' '
-# Maximum number of characters on a single line.
+# Maximum number of characters on a single line. Pylint's default of 100 is
+# based on PEP 8's guidance that teams may choose line lengths up to 99
+# characters.
max-line-length=150
# Maximum number of lines in a module.
-max-module-lines=2000
+max-module-lines=1000
# Allow the body of a class to be on the same line as the declaration if body
# contains single statement.
@@ -421,11 +436,21 @@ confidence=HIGH,
# --enable=similarities". If you want to run only the classes checker, but have
# no Warning level messages displayed, use "--disable=all --enable=classes
# --disable=W".
-disable=too-many-arguments,
- too-many-locals,
- too-many-branches,
- protected-access
-
+disable=raw-checker-failed,
+ bad-inline-option,
+ locally-disabled,
+ file-ignored,
+ suppressed-message,
+ useless-suppression,
+ deprecated-pragma,
+ use-symbolic-message-instead,
+ use-implicit-booleaness-not-comparison-to-string,
+ use-implicit-booleaness-not-comparison-to-zero,
+ C0116,
+ C0114,
+ W0621,
+ E0601,
+ W0718
# Enable the message, report, category or checker with the given id(s). You can
# either give multiple identifier separated by comma (,) or put this option
@@ -443,9 +468,13 @@ timeout-methods=requests.api.delete,requests.api.get,requests.api.head,requests.
[MISCELLANEOUS]
+# Whether or not to search for fixme's in docstrings.
+check-fixme-in-docstring=no
+
# List of note tags to take in consideration, separated by a comma.
notes=FIXME,
- XXX
+ XXX,
+ TODO
# Regular expression of note tags to take in consideration.
notes-rgx=
@@ -465,7 +494,7 @@ never-returning-functions=sys.exit,argparse.parse_error
# Let 'consider-using-join' be raised when the separator to join on would be
# non-empty (resulting in expected fixes of the type: ``"- " + " -
# ".join(items)``)
-# suggest-join-with-non-empty-separator=yes
+suggest-join-with-non-empty-separator=yes
[REPORTS]
@@ -481,10 +510,10 @@ evaluation=max(0, 0 if fatal else 10.0 - ((float(5 * error + warning + refactor
# used to format the message information. See doc for all details.
msg-template=
-# Set the output format. Available formats are: text, parseable, colorized,
-# json2 (improved json format), json (old json format) and msvs (visual
-# studio). You can also give a reporter class, e.g.
-# mypackage.mymodule.MyReporterClass.
+# Set the output format. Available formats are: 'text', 'parseable',
+# 'colorized', 'json2' (improved json format), 'json' (old json format), msvs
+# (visual studio) and 'github' (GitHub actions). You can also give a reporter
+# class, e.g. mypackage.mymodule.MyReporterClass.
#output-format=
# Tells whether to display a full report or only the messages.
@@ -586,7 +615,7 @@ ignored-classes=optparse.Values,thread._local,_thread._local,argparse.Namespace
# of finding the hint is based on edit distance.
missing-member-hint=yes
-# The minimum edit distance a name should have in order to be considered a
+# The maximum edit distance a name should have in order to be considered a
# similar match for a missing member name.
missing-member-hint-distance=1
@@ -630,4 +659,4 @@ init-import=no
# List of qualified module names which can have objects that can redefine
# builtins.
-redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io
\ No newline at end of file
+redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..261eeb9
--- /dev/null
+++ b/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.
diff --git a/README.md b/README.md
index 50cfc2f..5271aa3 100644
--- a/README.md
+++ b/README.md
@@ -27,17 +27,23 @@
+
+
## π° News
- **`2025/12/06`**: πππ Paper submitted on [Arxiv](https://arxiv.org/pdf/2512.06112).
+
+
## π
οΈ Roadmap
| Status | Milestone | ETA |
| :----: | :----------------------------------------------------------------------------------------------------: | :--------: |
-| π | **[Releasing the inference source code](https://github.com/fudan-generative-vision/WAM-Flow)** | 2025.12.21 |
-| π | **[Pretrained models on Huggingface](https://huggingface.co/fudan-generative-ai/WAM-Flow)** | TBD |
-| π | **[Releasing the training scripts](#training)** | TBD |
+| β
| **[Release the SFT and inference code](https://github.com/fudan-generative-vision/WAM-Flow)** | 2025.12.19 |
+| π | **[Pretrained models on Huggingface](https://huggingface.co/fudan-generative-ai/WAM-Flow)** | TBD |
+| π | **[Release the evaluation code](https://huggingface.co/fudan-generative-ai/WAM-Flow)** | TBD |
+| π | **[Release the RL code](https://github.com/fudan-generative-vision/WAM-Flow)** | TBD |
+| π | **[Release the pre-processed training data](#training)** | TBD |
## πΈ Showcase
@@ -45,14 +51,68 @@
## π Qualitative Results on NAVSIM
### NAVSIM-v1 benchmark results
-
+
+
+
+
+
",
+ stop_token_ids=[2],
+ )
+)
+
+# llama2 template
+# reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212
+register_conv_template(
+ Conversation(
+ name="llama-2",
+ system_template="[INST] <>\n{system_message}\n<>\n\n",
+ roles=("[INST]", "[/INST]"),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.LLAMA2,
+ sep=" ",
+ sep2=" ",
+ stop_token_ids=[2],
+ )
+)
+
+
+# deepseek template
+register_conv_template(
+ Conversation(
+ name="fudoki",
+ system_template="{system_message}",
+ # system_message="You are a helpful assistant. Please answer truthfully and write out your "
+ # "thinking step by step to be sure you get the right answer.",
+ system_message="",
+ roles=("User", "Assistant"),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.DeepSeek,
+ sep="\n\n",
+ sep2="<ο½endβofβsentenceο½>",
+ stop_token_ids=[100001],
+ stop_str=["User:", "<ο½endβofβsentenceο½>"],
+ )
+)
+
+register_conv_template(
+ Conversation(
+ name="plain",
+ system_template="",
+ system_message="",
+ roles=("", ""),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.PLAIN,
+ sep="",
+ sep2="",
+ stop_token_ids=[2],
+ stop_str=[""],
+ )
+)
+
+
+register_conv_template(
+ Conversation(
+ name="alignment",
+ system_template="",
+ system_message="",
+ roles=("", ""),
+ messages=(),
+ offset=0,
+ sep_style=SeparatorStyle.ALIGNMENT,
+ sep="",
+ sep2="",
+ stop_token_ids=[2],
+ stop_str=[""],
+ )
+)
+
+
+if __name__ == "__main__":
+ print("fudoki template:")
+ conv = get_conv_template("fudoki")
+ conv.append_message(conv.roles[0], "Hello!")
+ conv.append_message(conv.roles[1], "Hi! This is Tony.")
+ conv.append_message(conv.roles[0], "Who are you?")
+ conv.append_message(conv.roles[1], "I am a helpful assistant.")
+ conv.append_message(conv.roles[0], "How are you?")
+ conv.append_message(conv.roles[1], None)
+ print(conv.get_prompt())
diff --git a/fudoki/janus/utils/io.py b/fudoki/janus/utils/io.py
new file mode 100644
index 0000000..e8e9da7
--- /dev/null
+++ b/fudoki/janus/utils/io.py
@@ -0,0 +1,89 @@
+# Copyright (c) 2023-2024 DeepSeek.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of
+# this software and associated documentation files (the "Software"), to deal in
+# the Software without restriction, including without limitation the rights to
+# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+# the Software, and to permit persons to whom the Software is furnished to do so,
+# subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+import json
+from typing import Dict, List
+
+import PIL.Image
+import torch
+import base64
+import io
+from transformers import AutoModelForCausalLM
+
+from janus.models import MultiModalityCausalLM, VLChatProcessor
+
+
+def load_pretrained_model(model_path: str):
+ vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
+ tokenizer = vl_chat_processor.tokenizer
+
+ vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
+ model_path, trust_remote_code=True
+ )
+ vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
+
+ return tokenizer, vl_chat_processor, vl_gpt
+
+
+def load_pil_images(conversations: List[Dict[str, str]]) -> List[PIL.Image.Image]:
+ """
+
+ Support file path or base64 images.
+
+ Args:
+ conversations (List[Dict[str, str]]): the conversations with a list of messages. An example is :
+ [
+ {
+ "role": "User",
+ "content": "\nExtract all information from this image and convert them into markdown format.",
+ "images": ["./examples/table_datasets.png"]
+ },
+ {"role": "Assistant", "content": ""},
+ ]
+
+ Returns:
+ pil_images (List[PIL.Image.Image]): the list of PIL images.
+
+ """
+
+ pil_images = []
+
+ for message in conversations:
+ if "images" not in message:
+ continue
+
+ for image_data in message["images"]:
+ if image_data.startswith("data:image"):
+ # Image data is in base64 format
+ _, image_data = image_data.split(",", 1)
+ image_bytes = base64.b64decode(image_data)
+ pil_img = PIL.Image.open(io.BytesIO(image_bytes))
+ else:
+ # Image data is a file path
+ pil_img = PIL.Image.open(image_data)
+ pil_img = pil_img.convert("RGB")
+ pil_images.append(pil_img)
+
+ return pil_images
+
+
+def load_json(filepath):
+ with open(filepath, "r") as f:
+ data = json.load(f)
+ return data
diff --git a/fudoki/model.py b/fudoki/model.py
new file mode 100644
index 0000000..4cfe97e
--- /dev/null
+++ b/fudoki/model.py
@@ -0,0 +1,11 @@
+from transformers import AutoModelForCausalLM
+from fudoki.janus.models import MultiModalityCausalLM
+
+
+def instantiate_model(pretrained_weight_path):
+
+ vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
+ pretrained_weight_path, trust_remote_code=True
+ )
+ model = vl_gpt
+ return model
diff --git a/fudoki/vq_model.py b/fudoki/vq_model.py
new file mode 100644
index 0000000..4ee2b48
--- /dev/null
+++ b/fudoki/vq_model.py
@@ -0,0 +1,424 @@
+# Inherited from LlamaGen (https://arxiv.org/abs/2406.06525), which is modified from:
+# taming-transformers: https://github.com/CompVis/taming-transformers
+# maskgit: https://github.com/google-research/maskgit
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class ModelArgs:
+ codebook_size: int = 16384
+ codebook_embed_dim: int = 8
+ codebook_l2_norm: bool = True
+ codebook_show_usage: bool = True
+ commit_loss_beta: float = 0.25
+ entropy_loss_ratio: float = 0.0
+
+ encoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
+ decoder_ch_mult: List[int] = field(default_factory=lambda: [1, 1, 2, 2, 4])
+ z_channels: int = 256
+ dropout_p: float = 0.0
+
+
+
+class VQModel(nn.Module):
+ def __init__(self, config: ModelArgs):
+ super().__init__()
+ self.config = config
+ self.encoder = Encoder(ch_mult=config.encoder_ch_mult, z_channels=config.z_channels, dropout=config.dropout_p)
+ self.decoder = Decoder(ch_mult=config.decoder_ch_mult, z_channels=config.z_channels, dropout=config.dropout_p)
+
+ self.quantize = VectorQuantizer(config.codebook_size, config.codebook_embed_dim,
+ config.commit_loss_beta, config.entropy_loss_ratio,
+ config.codebook_l2_norm, config.codebook_show_usage)
+ self.quant_conv = nn.Conv2d(config.z_channels, config.codebook_embed_dim, 1)
+ self.post_quant_conv = nn.Conv2d(config.codebook_embed_dim, config.z_channels, 1)
+
+ def encode(self, x):
+ h = self.encoder(x)
+ h = self.quant_conv(h)
+ quant, emb_loss, info = self.quantize(h)
+ return quant, emb_loss, info
+
+ def decode(self, quant):
+ quant = self.post_quant_conv(quant)
+ dec = self.decoder(quant)
+ return dec
+
+ def decode_code(self, code_b, shape=None, channel_first=True):
+ quant_b = self.quantize.get_codebook_entry(code_b, shape, channel_first)
+ dec = self.decode(quant_b)
+ return dec
+
+ def forward(self, input):
+ quant, diff, _ = self.encode(input)
+ dec = self.decode(quant)
+ return dec, diff
+
+
+
+class Encoder(nn.Module):
+ def __init__(self, in_channels=3, ch=128, ch_mult=(1,1,2,2,4), num_res_blocks=2,
+ norm_type='group', dropout=0.0, resamp_with_conv=True, z_channels=256):
+ super().__init__()
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.conv_in = nn.Conv2d(in_channels, ch, kernel_size=3, stride=1, padding=1)
+
+ # downsampling
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.conv_blocks = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ conv_block = nn.Module()
+ # res & attn
+ res_block = nn.ModuleList()
+ attn_block = nn.ModuleList()
+ block_in = ch*in_ch_mult[i_level]
+ block_out = ch*ch_mult[i_level]
+ for _ in range(self.num_res_blocks):
+ res_block.append(ResnetBlock(block_in, block_out, dropout=dropout, norm_type=norm_type))
+ block_in = block_out
+ if i_level == self.num_resolutions - 1:
+ attn_block.append(AttnBlock(block_in, norm_type))
+ conv_block.res = res_block
+ conv_block.attn = attn_block
+ # downsample
+ if i_level != self.num_resolutions-1:
+ conv_block.downsample = Downsample(block_in, resamp_with_conv)
+ self.conv_blocks.append(conv_block)
+
+ # middle
+ self.mid = nn.ModuleList()
+ self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+ self.mid.append(AttnBlock(block_in, norm_type=norm_type))
+ self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+
+ # end
+ self.norm_out = Normalize(block_in, norm_type)
+ self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1)
+
+
+ def forward(self, x):
+ h = self.conv_in(x)
+ # downsampling
+ for i_level, block in enumerate(self.conv_blocks):
+ for i_block in range(self.num_res_blocks):
+ h = block.res[i_block](h)
+ if len(block.attn) > 0:
+ h = block.attn[i_block](h)
+ if i_level != self.num_resolutions - 1:
+ h = block.downsample(h)
+
+ # middle
+ for mid_block in self.mid:
+ h = mid_block(h)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+
+class Decoder(nn.Module):
+ def __init__(self, z_channels=256, ch=128, ch_mult=(1,1,2,2,4), num_res_blocks=2, norm_type="group",
+ dropout=0.0, resamp_with_conv=True, out_channels=3):
+ super().__init__()
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+
+ block_in = ch*ch_mult[self.num_resolutions-1]
+ # z to block_in
+ self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+ # middle
+ self.mid = nn.ModuleList()
+ self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+ self.mid.append(AttnBlock(block_in, norm_type=norm_type))
+ self.mid.append(ResnetBlock(block_in, block_in, dropout=dropout, norm_type=norm_type))
+
+ # upsampling
+ self.conv_blocks = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ conv_block = nn.Module()
+ # res & attn
+ res_block = nn.ModuleList()
+ attn_block = nn.ModuleList()
+ block_out = ch*ch_mult[i_level]
+ for _ in range(self.num_res_blocks + 1):
+ res_block.append(ResnetBlock(block_in, block_out, dropout=dropout, norm_type=norm_type))
+ block_in = block_out
+ if i_level == self.num_resolutions - 1:
+ attn_block.append(AttnBlock(block_in, norm_type))
+ conv_block.res = res_block
+ conv_block.attn = attn_block
+ # downsample
+ if i_level != 0:
+ conv_block.upsample = Upsample(block_in, resamp_with_conv)
+ self.conv_blocks.append(conv_block)
+
+ # end
+ self.norm_out = Normalize(block_in, norm_type)
+ self.conv_out = nn.Conv2d(block_in, out_channels, kernel_size=3, stride=1, padding=1)
+
+ @property
+ def last_layer(self):
+ return self.conv_out.weight
+
+ def forward(self, z):
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ for mid_block in self.mid:
+ h = mid_block(h)
+
+ # upsampling
+ for i_level, block in enumerate(self.conv_blocks):
+ for i_block in range(self.num_res_blocks + 1):
+ h = block.res[i_block](h)
+ if len(block.attn) > 0:
+ h = block.attn[i_block](h)
+ if i_level != self.num_resolutions - 1:
+ h = block.upsample(h)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class VectorQuantizer(nn.Module):
+ def __init__(self, n_e, e_dim, beta, entropy_loss_ratio, l2_norm, show_usage):
+ super().__init__()
+ self.n_e = n_e
+ self.e_dim = e_dim
+ self.beta = beta
+ self.entropy_loss_ratio = entropy_loss_ratio
+ self.l2_norm = l2_norm
+ self.show_usage = show_usage
+
+ self.embedding = nn.Embedding(self.n_e, self.e_dim)
+ self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+ if self.l2_norm:
+ self.embedding.weight.data = F.normalize(self.embedding.weight.data, p=2, dim=-1)
+ if self.show_usage:
+ self.register_buffer("codebook_used", nn.Parameter(torch.zeros(65536)))
+
+
+ def forward(self, z):
+ # reshape z -> (batch, height, width, channel) and flatten
+ z = torch.einsum('b c h w -> b h w c', z).contiguous()
+ z_flattened = z.view(-1, self.e_dim)
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+ if self.l2_norm:
+ z = F.normalize(z, p=2, dim=-1)
+ z_flattened = F.normalize(z_flattened, p=2, dim=-1)
+ embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+ else:
+ embedding = self.embedding.weight
+
+ d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+ torch.sum(embedding**2, dim=1) - 2 * \
+ torch.einsum('bd,dn->bn', z_flattened, torch.einsum('n d -> d n', embedding))
+
+ min_encoding_indices = torch.argmin(d, dim=1)
+ z_q = embedding[min_encoding_indices].view(z.shape)
+ perplexity = None
+ min_encodings = None
+ vq_loss = None
+ commit_loss = None
+ entropy_loss = None
+ codebook_usage = 0
+
+ if self.show_usage and self.training:
+ cur_len = min_encoding_indices.shape[0]
+ self.codebook_used[:-cur_len] = self.codebook_used[cur_len:].clone()
+ self.codebook_used[-cur_len:] = min_encoding_indices
+ codebook_usage = len(torch.unique(self.codebook_used)) / self.n_e
+
+ # compute loss for embedding
+ if self.training:
+ vq_loss = torch.mean((z_q - z.detach()) ** 2)
+ commit_loss = self.beta * torch.mean((z_q.detach() - z) ** 2)
+ entropy_loss = self.entropy_loss_ratio * compute_entropy_loss(-d)
+
+ # preserve gradients
+ z_q = z + (z_q - z).detach()
+
+ # reshape back to match original input shape
+ z_q = torch.einsum('b h w c -> b c h w', z_q)
+
+ return z_q, (vq_loss, commit_loss, entropy_loss, codebook_usage), (perplexity, min_encodings, min_encoding_indices)
+
+ def get_codebook_entry(self, indices, shape=None, channel_first=True):
+ # shape = (batch, channel, height, width) if channel_first else (batch, height, width, channel)
+ if self.l2_norm:
+ embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+ else:
+ embedding = self.embedding.weight
+ z_q = embedding[indices] # (b*h*w, c)
+
+ if shape is not None:
+ if channel_first:
+ z_q = z_q.reshape(shape[0], shape[2], shape[3], shape[1])
+ # reshape back to match original input shape
+ z_q = z_q.permute(0, 3, 1, 2).contiguous()
+ else:
+ z_q = z_q.view(shape)
+ return z_q
+
+
+class ResnetBlock(nn.Module):
+ def __init__(self, in_channels, out_channels=None, conv_shortcut=False, dropout=0.0, norm_type='group'):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels, norm_type)
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ self.norm2 = Normalize(out_channels, norm_type)
+ self.dropout = nn.Dropout(dropout)
+ self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+ else:
+ self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+ def forward(self, x):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+ return x+h
+
+
+class AttnBlock(nn.Module):
+ def __init__(self, in_channels, norm_type='group'):
+ super().__init__()
+ self.norm = Normalize(in_channels, norm_type)
+ self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+ self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b,c,h,w = q.shape
+ q = q.reshape(b,c,h*w)
+ q = q.permute(0,2,1) # b,hw,c
+ k = k.reshape(b,c,h*w) # b,c,hw
+ w_ = torch.bmm(q,k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+ w_ = w_ * (int(c)**(-0.5))
+ w_ = F.softmax(w_, dim=2)
+
+ # attend to values
+ v = v.reshape(b,c,h*w)
+ w_ = w_.permute(0,2,1) # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v,w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+ h_ = h_.reshape(b,c,h,w)
+
+ h_ = self.proj_out(h_)
+
+ return x+h_
+
+
+def nonlinearity(x):
+ # swish
+ return x*torch.sigmoid(x)
+
+
+def Normalize(in_channels, norm_type='group'):
+ assert norm_type in ['group', 'batch']
+ if norm_type == 'group':
+ return nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+ elif norm_type == 'batch':
+ return nn.SyncBatchNorm(in_channels)
+
+
+class Upsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+ def forward(self, x):
+ x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+ if self.with_conv:
+ x = self.conv(x)
+ return x
+
+
+class Downsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+ def forward(self, x):
+ if self.with_conv:
+ pad = (0,1,0,1)
+ x = F.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ else:
+ x = F.avg_pool2d(x, kernel_size=2, stride=2)
+ return x
+
+
+def compute_entropy_loss(affinity, loss_type="softmax", temperature=0.01):
+ flat_affinity = affinity.reshape(-1, affinity.shape[-1])
+ flat_affinity /= temperature
+ probs = F.softmax(flat_affinity, dim=-1)
+ log_probs = F.log_softmax(flat_affinity + 1e-5, dim=-1)
+ if loss_type == "softmax":
+ target_probs = probs
+ else:
+ raise ValueError("Entropy loss {} not supported".format(loss_type))
+ avg_probs = torch.mean(target_probs, dim=0)
+ avg_entropy = - torch.sum(avg_probs * torch.log(avg_probs + 1e-5))
+ sample_entropy = - torch.mean(torch.sum(target_probs * log_probs, dim=-1))
+ loss = sample_entropy - avg_entropy
+ return loss
+
+
+#################################################################################
+# VQ Model Configs #
+#################################################################################
+def VQ_8(**kwargs):
+ return VQModel(ModelArgs(encoder_ch_mult=[1, 2, 2, 4], decoder_ch_mult=[1, 2, 2, 4], **kwargs))
+
+def VQ_16(**kwargs):
+ return VQModel(ModelArgs(encoder_ch_mult=[1, 1, 2, 2, 4], decoder_ch_mult=[1, 1, 2, 2, 4], **kwargs))
+
+VQ_models = {'VQ-16': VQ_16, 'VQ-8': VQ_8}
\ No newline at end of file
diff --git a/infer.py b/infer.py
new file mode 100644
index 0000000..051d12e
--- /dev/null
+++ b/infer.py
@@ -0,0 +1,332 @@
+# LINT_ME
+import os
+import argparse
+
+import numpy as np
+from PIL import Image
+import torch
+import torch.distributed as dist
+from torchvision import transforms
+from torch.backends import cudnn
+from transformers import set_seed
+
+from flow_matching.data.navsim import resize_pad
+from flow_matching.path import MixtureDiscreteSoftmaxProbPath
+from flow_matching.solver import MixtureDiscreteSoftmaxEulerSolver
+from fudoki.eval_loop import CFGScaledModel
+from fudoki.janus.models import VLChatProcessor
+from fudoki.model import instantiate_model
+
+
+VOCABULARY_SIZE_TXT = 102400
+VOCABULARY_SIZE_IMG = 16384
+IMG_LEN = 576
+
+
+def parse_arguments():
+ parser = argparse.ArgumentParser(description="Run the script with custom arguments.")
+ parser.add_argument(
+ "--seed", type=int, default=999,
+ help="Random seed for reproducibility."
+ )
+ parser.add_argument(
+ "--batch_size", type=int, default=1,
+ help="Batch size for processing."
+ )
+ parser.add_argument(
+ "--checkpoint_path", type=str, required=True,
+ help="Path to the checkpoint directory."
+ )
+ parser.add_argument(
+ "--processor_path", type=str, required=True,
+ help="Path to the processor."
+ )
+ parser.add_argument(
+ "--text_embedding_path", type=str, required=True,
+ help="Path to the text embedding."
+ )
+ parser.add_argument(
+ "--image_embedding_path", type=str, required=True,
+ help="Path to the image embedding."
+ )
+ parser.add_argument(
+ "--discrete_fm_steps", type=int, default=5,
+ help="Inference steps for discrete flow matching"
+ )
+ parser.add_argument(
+ "--txt_max_length", type=int, default=500,
+ help="Text length maximum"
+ )
+ parser.add_argument(
+ "--image_paths", type=str, required=True,
+ help="Path to the input image."
+ )
+ parser.add_argument(
+ "--output_dir", type=str, required=False,
+ help="Directory to save the output files."
+ )
+ return parser.parse_args()
+
+
+def extract_number_pairs(input_str, k=None):
+ """
+ Extract number pairs from a string and form tuples of (float(x), float(y))
+
+ Parameters:
+ input_str (str): Input string in the format "4.78,-0.01,9.69,-0.01,..."
+ k (int, optional): Number of tuples to extract. If None, extract all possible tuples
+
+ Returns:
+ tuple: Contains two elements
+ - list: List of successfully extracted tuples
+ - list: List of error messages
+ """
+ pairs = []
+ errors = []
+
+ # Check if input is empty or not a string
+ if not input_str or not isinstance(input_str, str):
+ errors.append("Input is not a valid string")
+ return pairs, errors
+
+ # Split the string
+ elements = input_str.split(',')
+
+ # Process each element and attempt to convert to float
+ numbers = []
+ for index, elem in enumerate(elements):
+ # Remove possible whitespace characters
+ elem_clean = elem.strip()
+ if not elem_clean:
+ errors.append(f"Empty value at position {index}")
+ continue
+
+ try:
+ num = float(elem_clean)
+ numbers.append(num)
+ except ValueError:
+ errors.append(f"Value '{elem_clean}' at position {index} \
+ cannot be converted to a number")
+
+ # Calculate maximum possible pairs
+ max_possible = len(numbers) // 2
+
+ # Determine number of pairs to extract
+ if k is None:
+ # Extract all possible pairs
+ num_to_extract = max_possible
+ else:
+ # Ensure k is a positive integer
+ try:
+ k = int(k)
+ if k <= 0:
+ errors.append(f"k value {k} must be a positive integer")
+ return pairs, errors
+ num_to_extract = min(k, max_possible)
+ except (ValueError, TypeError):
+ errors.append(f"k value {k} is not a valid integer")
+ return pairs, errors
+
+ # Extract number pairs
+ for i in range(num_to_extract):
+ x = numbers[2*i]
+ y = numbers[2*i + 1]
+ pairs.append((x, y))
+
+ # Check for unpaired numbers
+ if len(numbers) % 2 != 0:
+ errors.append(f"There are {len(numbers) % 2} unpaired number(s)")
+
+ # Check if requested k value was achieved
+ if k is not None and num_to_extract < k:
+ errors.append(f"Only {num_to_extract} valid number pairs can be extracted, \
+ less than the requested {k}")
+
+ return pairs, errors
+
+
+def main():
+ args = parse_arguments()
+
+ dist.init_process_group(
+ "nccl",
+ rank=int(os.environ["RANK"]),
+ world_size=int(os.environ["WORLD_SIZE"]),
+ )
+ local_rank = int(os.environ["LOCAL_RANK"])
+ torch.cuda.set_device(local_rank)
+ set_seed(args.seed)
+ cudnn.benchmark = True
+ device = torch.device(f"cuda:{local_rank}")
+
+ image_paths = args.image_paths.split(',')
+
+ vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(args.processor_path)
+ num_tokens_length = 0
+ num_tokens = [f"{x:.2f}" for x in np.linspace(-100, 100, 20001)]
+ num_tokens_length = len(num_tokens)
+ vl_chat_processor.tokenizer.add_tokens(num_tokens)
+
+ model = instantiate_model(args.checkpoint_path).to(device, dtype=torch.float32)
+ model.train(False)
+
+ batch_size = args.batch_size
+ discrete_fm_steps = args.discrete_fm_steps
+ txt_max_length = args.txt_max_length
+
+ cfg_weighted_model = CFGScaledModel(model=model, g_or_u='understanding')
+ with torch.no_grad():
+ path_txt = MixtureDiscreteSoftmaxProbPath(
+ mode='text',
+ embedding_path=args.text_embedding_path
+ )
+ path_txt.set_embedding(model.language_model.get_input_embeddings())
+ path_img = MixtureDiscreteSoftmaxProbPath(
+ mode='image',
+ embedding_path=args.image_embedding_path
+ )
+ solver = MixtureDiscreteSoftmaxEulerSolver(
+ model=cfg_weighted_model,
+ path_txt=path_txt,
+ path_img=path_img,
+ vocabulary_size_txt=VOCABULARY_SIZE_TXT + num_tokens_length,
+ vocabulary_size_img=VOCABULARY_SIZE_IMG,
+ )
+
+ imgs = []
+ if isinstance(image_paths, str):
+ image_paths = [image_paths]
+ for path in image_paths:
+ img = Image.open(path).convert("RGB")
+ transform = transforms.Compose([
+ transforms.Lambda(resize_pad),
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+ ])
+ imgs.append(transform(img))
+
+ if len(imgs) > 0:
+ imgs = torch.stack(imgs, dim=0) # [N, C, H, W]
+ img_len = len(imgs) * IMG_LEN
+ else:
+ imgs = None
+ img_len = IMG_LEN # default
+
+ generation_understanding_indicator = 0 # this is an understanding sample
+ conversation = [
+ {
+ "role": "User",
+ "content": (
+ "Here is front views of a driving vehicle:\n\n"
+ "The navigation information is: straight\n"
+ "The current position is (0.00,0.00)\n"
+ "Current velocity is: (8.34,0.18) and current accelerate is: (-0.83,0.28)\n"
+ "Predict the optimal driving action for the next 4 seconds with 8 new waypoints."
+ )
+ },
+ {
+ "role": "Assistant",
+ "content": ""
+ } # "3.88,-0.06,7.50,-0.07,10.86,-0.10,13.95,-0.11,16.75,-0.13,19.29,-0.15,21.60,-0.12,23.67,-0.11"
+ ]
+ sft_format = vl_chat_processor.apply_sft_template_for_multi_turn_prompts(
+ conversations=conversation,
+ sft_format=vl_chat_processor.sft_format,
+ system_prompt=vl_chat_processor.system_prompt,
+ )
+
+ # tokenize
+ input_ids = vl_chat_processor.tokenizer.encode(sft_format)
+ input_ids = torch.LongTensor(input_ids)
+ # add image tokens to the input_ids
+ image_token_mask = input_ids == vl_chat_processor.image_id
+ image_indices = image_token_mask.nonzero()
+ assert len(image_indices) == len(image_paths), \
+ f"Number of images ({len(image_paths)}) \
+ does not match the number of image tokens ({len(image_indices)})"
+
+ input_ids, _ = vl_chat_processor.add_image_token(
+ image_indices=image_indices,
+ input_ids=input_ids,
+ )
+
+ # pad tokens
+ original_input_id_len = input_ids.shape[0]
+ if original_input_id_len >= txt_max_length + img_len:
+ raise ValueError("Sentences too long, not supported so far...")
+
+ rows_to_pad = txt_max_length + img_len - input_ids.shape[0]
+ input_ids = torch.concat([
+ input_ids,
+ torch.LongTensor([vl_chat_processor.pad_id]).repeat(rows_to_pad)
+ ], dim=0)
+ attention_mask = torch.zeros((input_ids.shape[0]), dtype=torch.bool)
+ attention_mask[:] = True
+
+ # obtain image token mask and fill in img token_ids
+ if imgs is not None:
+ image_expanded_token_mask = (input_ids == vl_chat_processor.image_id).to(dtype=int)
+ image_expanded_mask_indices = torch.where(image_expanded_token_mask == 1)[0]
+ input_ids[image_expanded_mask_indices] = 0
+ else:
+ image_expanded_token_mask = torch.zeros_like(input_ids)
+
+ # obtain text token mask
+ # We assume that there is only one turn for assistant to respond
+ text_expanded_token_mask = torch.zeros_like(image_expanded_token_mask)
+ split_token = vl_chat_processor.tokenizer.encode("Assistant:", add_special_tokens=False)
+ split_token_length = len(split_token)
+
+ start_index = -1
+ for j in range(len(input_ids) - split_token_length + 1):
+ if input_ids[j:j + split_token_length].numpy().tolist() == split_token:
+ start_index = j
+ break
+ if start_index != -1:
+ text_expanded_token_mask[(start_index+split_token_length):] = 1
+ else:
+ raise ValueError("Split token not found in input_ids")
+
+ generation_or_understanding_mask = generation_understanding_indicator
+ data_info = {}
+ data_info['text_token_mask'] = text_expanded_token_mask.unsqueeze(0).repeat(batch_size, 1).to(device)
+ data_info['image_token_mask'] = image_expanded_token_mask.unsqueeze(0).repeat(batch_size, 1).to(device)
+ data_info['generation_or_understanding_mask'] = \
+ torch.Tensor([generation_or_understanding_mask]).unsqueeze(0).repeat(batch_size, 1).to(device).to(dtype=int)
+
+ data_info['attention_mask'] = attention_mask.unsqueeze(0).repeat(batch_size, 1).to(device)
+ data_info['sft_format'] = sft_format
+ if imgs is not None:
+ data_info['understanding_img'] = imgs.unsqueeze(0).to(device, dtype=torch.float32).repeat(batch_size, 1, 1, 1, 1)
+ data_info['has_understanding_img'] = torch.Tensor([True]).to(dtype=int).unsqueeze(0).repeat(batch_size, 1).to(device)
+ else:
+ data_info['understanding_img'] = torch.zeros((3, 384, 384)).unsqueeze(0).repeat(batch_size, 1, 1, 1).to(device)
+ data_info['has_understanding_img'] = torch.Tensor([False]).to(dtype=int).unsqueeze(0).repeat(batch_size, 1).to(device)
+ input_ids = torch.LongTensor(input_ids).unsqueeze(0).repeat(batch_size, 1).to(device)
+
+
+ x_0_txt = torch.randint(VOCABULARY_SIZE_TXT + num_tokens_length, input_ids.shape, dtype=torch.long, device=device)
+ x_init = x_0_txt * data_info['text_token_mask'] + input_ids * (1 - data_info['text_token_mask'])
+
+ synthetic_samples = solver.sample(
+ x_init=x_init,
+ step_size=1.0/discrete_fm_steps,
+ verbose=True,
+ return_intermediates=False,
+ div_free=0,
+ dtype_categorical=torch.float32,
+ datainfo=data_info,
+ cfg_scale=0,
+ )
+ sentence = vl_chat_processor.tokenizer.batch_decode(
+ synthetic_samples,
+ skip_special_tokens=True
+ )[0]
+ print("Sentence:", sentence)
+
+ waypoint = extract_number_pairs(sentence, k=8)[0]
+ print("Waypoint: ", waypoint)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..dab4802
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,17 @@
+omegaconf
+transformers==4.42.4
+timm
+tokenizers
+attrdict
+torch==2.0.1
+torchvision
+accelerate
+sentencepiece
+einops
+torchdiffeq
+matplotlib
+numpy==1.26.4
+icecream
+xformers==0.0.22
+diffusers==0.32.2
+deepspeed
\ No newline at end of file
diff --git a/script/infer.sh b/script/infer.sh
new file mode 100644
index 0000000..114b5f4
--- /dev/null
+++ b/script/infer.sh
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+CKPT_PATH="pretrained_model/wam-flow/navsim"
+FUDOKI_PATH="pretrained_model/fudoki"
+IMAGE_PATH="data/navsim_data/sensor_blobs/test/2021.09.09.17.18.51_veh-48_00889_01147/CAM_F0/9a6f0331d98258a0.jpg"
+
+torchrun --nproc_per_node 1 infer.py \
+ --checkpoint_path $CKPT_PATH \
+ --image_path $IMAGE_PATH \
+ --processor_path $FUDOKI_PATH \
+ --text_embedding_path $FUDOKI_PATH/text_embedding.pt \
+ --image_embedding_path $FUDOKI_PATH/image_embedding.pt \
+ --discrete_fm_steps 2 \
+ --seed 123
\ No newline at end of file
diff --git a/script/sft_debug.sh b/script/sft_debug.sh
new file mode 100644
index 0000000..fd310a6
--- /dev/null
+++ b/script/sft_debug.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+NUM_NODES=1
+NUM_GPUS=1
+
+config=config/debug.yaml
+output_dir=output/train/debug
+
+accelerate launch \
+ --config_file ./config/accelerate_config_ds2.yaml \
+ --machine_rank 0 \
+ --main_process_port 12345 \
+ --num_machines $NUM_NODES \
+ --num_processes $NUM_GPUS \
+ train.py \
+ --config $config \
+ --output_dir $output_dir
diff --git a/script/sft_navsim.sh b/script/sft_navsim.sh
new file mode 100644
index 0000000..aedda23
--- /dev/null
+++ b/script/sft_navsim.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+NUM_NODES=1
+NUM_GPUS=1
+
+config=config/sft_navsim.yaml
+output_dir=output/train/debug
+
+accelerate launch \
+ --config_file ./config/accelerate_config_ds2.yaml \
+ --machine_rank 0 \
+ --main_process_port 12345 \
+ --num_machines $NUM_NODES \
+ --num_processes $NUM_GPUS \
+ train.py \
+ --config $config \
+ --output_dir $output_dir
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..6c8398f
--- /dev/null
+++ b/train.py
@@ -0,0 +1,451 @@
+# LINT_ME
+import argparse
+import os
+import logging
+import shutil
+import copy
+import random
+
+import numpy as np
+import transformers
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torch.nn import CrossEntropyLoss
+import diffusers
+from diffusers.optimization import get_scheduler
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from omegaconf import OmegaConf
+from tqdm.auto import tqdm
+
+from fudoki.model import instantiate_model
+from fudoki.janus.models import VLChatProcessor
+from flow_matching.path import MixtureDiscreteSoftmaxProbPath
+from flow_matching.data.navsim import SupervisedDataset
+from flow_matching.utils.flow import get_source_distribution
+
+
+logger = get_logger(__name__, log_level="INFO")
+
+
+@torch.no_grad()
+def init_numeric_and_special_tokens(
+ model,
+ tokenizer,
+ numeric_tokens,
+ noise_scale: float = 0.01,
+):
+ emb = model.get_input_embeddings().weight
+ device = emb.device
+ dim = emb.shape[1]
+
+ def tok_ids_for_text(text: str):
+ # Get subword ids (no specials), filter out -100/None if any
+ ids = tokenizer.encode(text,add_special_tokens=False)
+ return [i for i in ids if isinstance(i, int) and i >= 0]
+
+ # ---- Build a numeric "base" vector from digits and dot ----
+ digit_ids = []
+ for d in "0123456789":
+ _ids = tok_ids_for_text(d)
+ digit_ids.extend(_ids)
+ dot_ids = tok_ids_for_text(".")
+
+ base_chunks = []
+ if digit_ids:
+ base_chunks.append(emb[torch.tensor(digit_ids, device=device)].mean(dim=0))
+ if dot_ids:
+ base_chunks.append(emb[torch.tensor(dot_ids, device=device)].mean(dim=0))
+ if base_chunks:
+ numeric_base = torch.stack(base_chunks, dim=0).mean(dim=0)
+ else:
+ # fallback if tokenizer lacks digits/dot as standalone pieces
+ numeric_base = torch.zeros(dim, device=device)
+
+ # ---- Initialize numeric tokens ----
+ for t in numeric_tokens:
+ tid = tokenizer.convert_tokens_to_ids(t)
+ if tid is None or tid < 0:
+ continue
+ noise = noise_scale * torch.randn(dim, device=device)
+ emb[tid] = numeric_base + noise
+
+
+def training_step(
+ model,
+ x_1,
+ source_distribution,
+ data_info,
+ path,
+ time_epsilon = 0.001,
+ loss_fn = CrossEntropyLoss(),
+ stage="s1",
+ vl_chat_processor=None,
+ args=None,
+):
+ x_0 = source_distribution.sample_like(x_1)
+ t = torch.rand(x_1.shape[0], device=x_1.device) * (1.0 - time_epsilon)
+
+ if stage == "s1":
+ x_t = x_1
+ elif stage == "s2":
+ # update emb layer when using num tokenizer
+ path_sample = path.sample(x_0, x_1, t)
+ x_t = path_sample.x_t
+
+ # text_token_mask==1 ==> generated text token
+ x_t = x_t * data_info['text_token_mask'] + x_1 * (1 - data_info['text_token_mask'])
+ data_info['understanding_img'] = data_info['understanding_img'].to(dtype=model.dtype)
+
+ _, txt_logits = model(x_t, data_info)
+
+ b, _, c = txt_logits.shape
+ mask = data_info['text_token_mask'].unsqueeze(-1).bool()
+ txt_logits = txt_logits.masked_select(mask)
+ txt_logits = txt_logits.view(b, -1, c)
+ x_1 = x_1.masked_select(mask.squeeze(-1)).view(b, -1)
+
+ loss = ce_loss = loss_fn(txt_logits.flatten(0, 1), x_1.flatten(0, 1)).mean()
+ loss_dict = {"ce_loss": ce_loss.detach().item()}
+
+ if stage == "s2":
+ start_mask = x_1 >= vl_chat_processor.num_start_id
+ end_mask = x_1 <= vl_chat_processor.num_end_id - 1
+ action_mask = start_mask & end_mask
+
+ if action_mask.any():
+ if args.l2_loss_weight > 0:
+ pred_probabilities = F.softmax(txt_logits, dim=-1)
+ pred_ids = torch.argmax(pred_probabilities, dim=-1)
+ pred_num_ids = pred_ids.masked_select(action_mask)
+ pred_nums = vl_chat_processor.min_num + (pred_num_ids - vl_chat_processor.num_start_id) * vl_chat_processor.interval
+ pred_nums = torch.clip(pred_nums, vl_chat_processor.min_num, vl_chat_processor.max_num)
+
+ tgt_num_ids = x_1.masked_select(action_mask) # [N]
+ tgt_nums = vl_chat_processor.min_num + (tgt_num_ids - vl_chat_processor.num_start_id) * vl_chat_processor.interval
+
+ l2_loss = args.l2_loss_weight + F.mse_loss(pred_nums, tgt_nums, reduction="mean")
+ loss = loss + l2_loss
+ loss_dict["l2_loss"] = l2_loss.detach().item()
+
+ loss_dict["loss"] = loss.detach().item()
+ return loss, loss_dict
+
+
+def main(args):
+ accelerator = Accelerator(
+ mixed_precision=args.mixed_precision,
+ gradient_accumulation_steps=args.accumulate_grad_batches,
+ log_with="tensorboard",
+ project_dir=args.output_dir
+ )
+
+ logging.basicConfig(
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+ datefmt="%m/%d/%Y %H:%M:%S",
+ level=logging.INFO,
+ )
+
+ logger.info(accelerator.state, main_process_only=False)
+ if accelerator.is_local_main_process:
+ transformers.utils.logging.set_verbosity_warning()
+ diffusers.utils.logging.set_verbosity_info()
+ else:
+ transformers.utils.logging.set_verbosity_error()
+ diffusers.utils.logging.set_verbosity_error()
+
+ # random seed
+ seed = args.seed + accelerator.process_index
+ if args.random_seed:
+ seed = seed + random.randint(0, 500)
+ set_seed(seed)
+ logger.info(f"accelerator.process_index: {accelerator.process_index}, seed: {seed} \n")
+
+ # work dir
+ if accelerator.is_main_process:
+ os.makedirs(args.output_dir, exist_ok=True)
+ config_path = os.path.join(args.output_dir, "config.yaml")
+ OmegaConf.save(args, config_path)
+ accelerate_config_path = os.path.join(args.output_dir, "accelerate_config_ds2.yaml")
+ shutil.copyfile(
+ "./config/accelerate_config_ds2.yaml",
+ accelerate_config_path
+ )
+ accelerator.wait_for_everyone()
+
+ # dtype
+ weight_dtype = torch.float32
+ if accelerator.mixed_precision == "fp16":
+ weight_dtype = torch.float16
+ elif accelerator.mixed_precision == "bf16":
+ weight_dtype = torch.bfloat16
+
+ # prepare dataset
+ vl_chat_processor = VLChatProcessor.from_pretrained(args.model_path)
+ if args.use_quantize:
+ origin_len = len(vl_chat_processor.tokenizer)
+ num_tokens = [f"{x:.2f}" for x in np.linspace(-100, 100, 20001)]
+ num_tokens_length = len(num_tokens)
+ vl_chat_processor.tokenizer.add_tokens(num_tokens)
+
+ vl_chat_processor.num_start_id = origin_len
+ vl_chat_processor.num_end_id = origin_len + num_tokens_length - 1
+ vl_chat_processor.min_num = -100
+ vl_chat_processor.max_num = 100
+ vl_chat_processor.interval = 0.01
+
+ logger.info(f"Total number tokens: {num_tokens_length}")
+
+ # data
+ dataset = SupervisedDataset(
+ data_list=args.data_list,
+ vl_chat_processor=vl_chat_processor,
+ txt_max_length=args.txt_max_length
+ )
+ dataloader = DataLoader(
+ dataset,
+ shuffle=True,
+ batch_size=args.batch_size,
+ num_workers=args.dataloader_num_workers,
+ pin_memory=True
+ )
+ logger.info(f"Max txt length: {args.txt_max_length}")
+ logger.info(f"Total data samples: {len(dataset)}")
+
+ # prepare model
+ stage = args.stage
+ logger.info(f"Training stege: {stage}")
+ model = instantiate_model(
+ args.pretrain_model_path
+ ).to(weight_dtype)
+ model.uncond_prob = args.uncond_prob
+
+ if os.path.exists(args.pretrain_path):
+ sd = torch.load(args.pretrain_path, map_location='cpu')
+ model.load_state_dict(sd, strict=True)
+ model = model.to(weight_dtype)
+ logger.info(f"Loading pretrain ckpt from {args.pretrain_path}")
+
+ if stage == "s1":
+ model.language_model.resize_token_embeddings(args.vocab_size + num_tokens_length)
+ init_numeric_and_special_tokens(
+ model.language_model,
+ vl_chat_processor.tokenizer,
+ numeric_tokens=num_tokens
+ )
+ elif stage == "s2":
+ if os.path.exists(args.new_embedding_path):
+ old_emb = copy.deepcopy(model.language_model.get_input_embeddings())
+
+ with torch.serialization.safe_globals([torch.nn.modules.sparse.Embedding]):
+ new_emb_state = torch.load(args.new_embedding_path, map_location="cpu")
+ logger.info(f"Loading new embedding from {args.new_embedding_path}")
+
+ if isinstance(new_emb_state, dict) and "weight" in new_emb_state:
+ weight = new_emb_state["weight"]
+ new_emb = torch.nn.Embedding(weight.size(0), weight.size(1))
+ new_emb.load_state_dict(new_emb_state)
+ else:
+ new_emb = new_emb_state
+
+ model.language_model.resize_token_embeddings(args.vocab_size + num_tokens_length)
+
+ # origin_len = old_emb.weight.shape[0]
+ origin_len = vl_chat_processor.num_start_id
+ new_emb.weight.data[:origin_len, :] = old_emb.weight.data[:origin_len, :]
+
+ model.language_model.set_input_embeddings(new_emb)
+
+ if os.path.exists(args.ckpt_path):
+ if model.language_model.model.embed_tokens.weight.shape[0] != args.vocab_size + num_tokens_length:
+ model.language_model.resize_token_embeddings(args.vocab_size + num_tokens_length)
+ sd = torch.load(args.ckpt_path, map_location='cpu')
+ model.load_state_dict(sd, strict=True)
+ model = model.to(weight_dtype)
+ logger.info(f"Loading ckpt from {args.ckpt_path}")
+
+ # prepare path
+ path = MixtureDiscreteSoftmaxProbPath(
+ mode='text',
+ embedding_path=args.text_embedding_path
+ )
+ if args.use_quantize:
+ path.set_embedding(model.language_model.get_input_embeddings())
+ else:
+ logger.info("No quantize!")
+
+ logger.info(f"path.a = {path.a}")
+ logger.info(f"path.c = {path.c}")
+
+ # set trainable params
+ model.requires_grad_(False)
+ if stage == "s1":
+ model.language_model.requires_grad_(False)
+ model.language_model.model.embed_tokens.requires_grad_(True)
+ model.language_model.lm_head.requires_grad_(True)
+ elif stage == "s2":
+ model.language_model.requires_grad_(True)
+ if args.train_llm_emb:
+ model.language_model.model.embed_tokens.requires_grad_(True)
+ else:
+ model.language_model.model.embed_tokens.requires_grad_(False)
+
+ trainable_params = list(
+ filter(lambda p: p.requires_grad, model.parameters())
+ )
+
+ # log trainable params
+ # for name, param in model.named_parameters():
+ # if param.requires_grad:
+ # logger.info(f"Trainable params: {name}")
+
+ num_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ logger.info(f"Total trainable parameters: {num_trainable/1e9:.3} B")
+
+ optimizer = torch.optim.AdamW(
+ trainable_params,
+ lr=args.learning_rate,
+ betas=(0.9, 0.95),
+ weight_decay=0.05,
+ )
+
+ # lr scheduler
+ lr_scheduler = get_scheduler(
+ args.lr_scheduler_type,
+ optimizer=optimizer,
+ num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
+ num_training_steps=args.max_train_steps * accelerator.num_processes
+ )
+
+ # accelerator
+ model, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+ model, optimizer, dataloader, lr_scheduler
+ )
+
+ source_distribution = get_source_distribution(
+ source_distribution=args.source_distribution,
+ vocab_size=args.vocab_size + num_tokens_length if args.use_quantize else args.vocab_size,
+ )
+
+
+ global_step = 0
+ initial_global_step = 0
+
+ progress_bar = tqdm(
+ range(0, args.max_train_steps),
+ initial=initial_global_step,
+ desc="Steps",
+ disable=not accelerator.is_local_main_process,
+ )
+
+ # training loop
+ for epoch in range(args.max_epochs):
+ logger.info(f"Epoch {epoch + 1}/{args.max_epochs}")
+ logger.info(f"training sample length: {len(dataloader)}")
+
+ for _, batch in enumerate(dataloader):
+ with accelerator.accumulate(model):
+ x_1 = batch["input_ids"].to(dtype=torch.long)
+ loss, logs = training_step(
+ x_1=x_1,
+ model=model,
+ source_distribution=source_distribution,
+ data_info=batch,
+ path=path,
+ stage=stage,
+ vl_chat_processor=vl_chat_processor,
+ args=args
+ )
+ accelerator.backward(loss)
+
+ if accelerator.sync_gradients:
+ accelerator.clip_grad_norm_(
+ trainable_params,
+ args.max_grad_norm,
+ )
+ optimizer.step()
+ lr_scheduler.step()
+ optimizer.zero_grad()
+
+ if accelerator.sync_gradients:
+ progress_bar.update(1)
+ global_step += 1
+
+ logs["lr"] = optimizer.param_groups[0]['lr']
+
+ progress_bar.set_postfix(**logs)
+ accelerator.log(logs, step=global_step)
+
+ if global_step % args.checkpointing_steps == 0 \
+ and accelerator.is_main_process and args.checkpoints_total_limit is not None:
+
+ checkpoints = os.listdir(args.output_dir)
+ checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+ checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+ if len(checkpoints) >= args.checkpoints_total_limit:
+ num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+ removing_checkpoints = checkpoints[0:num_to_remove]
+
+ logger.info(
+ f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+ )
+ logger.info(f"Removing checkpoints: {', '.join(removing_checkpoints)}")
+
+ for removing_checkpoint in removing_checkpoints:
+ removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+ if os.path.exists(removing_checkpoint):
+ shutil.rmtree(removing_checkpoint)
+
+ accelerator.wait_for_everyone()
+ unwrap_net = accelerator.unwrap_model(model)
+ save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
+ if accelerator.is_main_process:
+ unwrap_net.save_pretrained(save_path, max_shard_size="20GB")
+ logger.info(f"Saved state to {save_path}")
+
+ if global_step >= args.max_train_steps:
+ break
+
+ if global_step >= args.max_train_steps:
+ break
+
+ accelerator.wait_for_everyone()
+ accelerator.end_training()
+ logger.info("training completed!")
+
+
+def parse_args():
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument("--config", type=str, required=True)
+ parser.add_argument("--output_dir", type=str, required=True)
+ parser.add_argument("--output_obs_dir", type=str, default=None)
+
+ return parser.parse_args()
+
+
+if __name__ == '__main__':
+ args = parse_args()
+
+ config = OmegaConf.load(args.config)
+
+ # merge args
+ args_dict = vars(args).copy()
+ args_dict.pop("config", None)
+ config_keys = set(config.keys())
+ cli_keys = set(args_dict.keys())
+
+ # check conflict
+ conflict_keys = cli_keys & config_keys
+ if conflict_keys:
+ print(f"Args conflict: {conflict_keys}")
+
+ # merge
+ merged_config = OmegaConf.merge(OmegaConf.create(args_dict), config)
+ args = merged_config
+
+ # training
+ main(args)