FLUX: Add support for setup configuration to publish module

flux/flux/__init__.py → flux/mlx_flux/__init__.py

@@ -1,5 +1,6 @@
 # Copyright © 2024 Apple Inc.
 
+from ._version import __version__
 from .datasets import Dataset, load_dataset
 from .flux import FluxPipeline
 from .lora import LoRALinear

flux/mlx_flux/_version.py

@@ -0,0 +1,3 @@
+# Copyright © 2023-2024 Apple Inc.
+
+__version__ = "0.1.0"

flux/dreambooth.py → flux/mlx_flux/dreambooth.py

@@ -13,7 +13,9 @@
 from mlx.utils import tree_flatten, tree_map, tree_reduce
 from PIL import Image
 
-from flux import FluxPipeline, Trainer, load_dataset, save_config
+from .datasets import load_dataset
+from .flux import FluxPipeline
+from .trainer import Trainer
 
 
 def generate_progress_images(iteration, flux, args):
@@ -153,7 +155,7 @@ def setup_arg_parser():
     return parser
 
 
-if __name__ == "__main__":
+def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
@@ -290,3 +292,7 @@ def step(x, t5_feat, clip_feat, guidance, prev_grads, perform_step):
 
     save_adapters("final_adapters.safetensors", flux, args)
     print(f"Training successful. Saved final weights to {args.adapter_file}.")
+
+
+if __name__ == "__main__":
+    main()

flux/txt2image.py → flux/mlx_flux/txt2image.py

@@ -8,7 +8,7 @@
 from PIL import Image
 from tqdm import tqdm
 
-from flux import FluxPipeline
+from .flux import FluxPipeline
 
 
 def to_latent_size(image_size):
@@ -39,7 +39,7 @@ def load_adapter(flux, adapter_file, fuse=False):
         flux.fuse_lora_layers()
 
 
-if __name__ == "__main__":
+def build_parser():
     parser = argparse.ArgumentParser(
         description="Generate images from a textual prompt using stable diffusion"
     )
@@ -62,7 +62,11 @@ def load_adapter(flux, adapter_file, fuse=False):
     parser.add_argument("--adapter")
     parser.add_argument("--fuse-adapter", action="store_true")
     parser.add_argument("--no-t5-padding", dest="t5_padding", action="store_false")
-    args = parser.parse_args()
+    return parser
+
+
+def main():
+    args = build_parser().parse_args()
 
     # Load the models
     flux = FluxPipeline("flux-" + args.model, t5_padding=args.t5_padding)
@@ -148,3 +152,7 @@ def load_adapter(flux, adapter_file, fuse=False):
         print(f"Peak memory used for the generation: {peak_mem_generation:.3f}GB")
         print(f"Peak memory used for the decoding:   {peak_mem_decoding:.3f}GB")
         print(f"Peak memory used overall:            {peak_mem_overall:.3f}GB")
+
+
+if __name__ == "__main__":
+    main()

flux/setup.py

@@ -0,0 +1,61 @@
+# Copyright © 2024 Apple Inc.
+
+import os
+import sys
+from pathlib import Path
+
+from setuptools import find_namespace_packages, setup
+
+# 获取当前文件的父目录（项目根目录）
+ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+package_dir = os.path.join(ROOT_DIR, "mlx_flux")
+
+# 定义依赖列表
+requirements = []
+if os.path.exists(os.path.join(ROOT_DIR, "requirements.txt")):
+    with open(os.path.join(ROOT_DIR, "requirements.txt")) as fid:
+        requirements = [l.strip() for l in fid.readlines() if l.strip()]
+
+# 添加包路径
+sys.path.append(package_dir)
+
+from _version import __version__
+
+try:
+    with open(os.path.join(ROOT_DIR, "README.md"), encoding="utf-8") as f:
+        long_description = f.read()
+except FileNotFoundError:
+    long_description = "FLUX.1 on Apple silicon with MLX and the Hugging Face Hub"
+
+setup(
+    name="mlx-flux",
+    version=__version__,
+    description="FLUX.1 on Apple silicon with MLX and the Hugging Face Hub",
+    long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type="text/markdown",
+    readme="README.md",
+    author_email="[email protected]",
+    author="MLX Contributors",
+    url="https://github.com/ml-explore/mlx-examples",
+    license="MIT",
+    install_requires=requirements,
+    # Package configuration
+    packages=find_namespace_packages(
+        include=["mlx_flux", "mlx_flux.*"]
+    ),  # 明确指定包含的包
+    package_data={
+        "mlx_flux": ["*.py"],
+    },
+    include_package_data=True,
+    python_requires=">=3.8",
+    entry_points={
+        "console_scripts": [
+            # generate images
+            "mlx_flux.generate = mlx_flux.txt2image:main",
+            "mlx_flux.txt2image = mlx_flux.txt2image:main",
+            # fine-tuning model
+            "mlx_flux.lora = mlx_flux.dreambooth:main",
+            "mlx_flux.dreambooth = mlx_flux.dreambooth:main",
+        ]
+    },
+)