feat: add support for uv and ruff, blazing-fast tools written in Rust

.python-version

@@ -0,0 +1 @@
+3.12

Dockerfile

@@ -9,6 +9,8 @@ RUN apt-get update && apt-get install -y \
 
 WORKDIR /app
 
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
 RUN curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -o ~/miniconda.sh \
     && sh ~/miniconda.sh -b -p /opt/conda \
     && rm ~/miniconda.sh
@@ -23,4 +25,4 @@ COPY tools tools
 COPY environment_setup.sh environment_setup.sh
 RUN ./environment_setup.sh
 
-CMD ["python", "-u", "-W", "ignore", "app/app_sana.py", "--share", "--config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml", "--model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth"]
+CMD ["uv", "run", "app/app_sana.py", "--share", "--config=configs/sana_config/1024ms/Sana_1600M_img1024.yaml", "--model_path=hf://Efficient-Large-Model/Sana_1600M_1024px/checkpoints/Sana_1600M_1024px.pth"]

app/app_sana.py

@@ -33,8 +33,8 @@
 from torchvision.utils import make_grid, save_image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from app import safety_check
-from app.sana_pipeline import SanaPipeline
+import safety_check
+from sana_pipeline import SanaPipeline
 
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "1") == "1"

environment_setup.sh

@@ -1,30 +1,30 @@
 #!/usr/bin/env bash
 set -e
 
-CONDA_ENV=${1:-""}
-if [ -n "$CONDA_ENV" ]; then
-    # This is required to activate conda environment
-    eval "$(conda shell.bash hook)"
+# Default virtual environment directory
+VENV_DIR=".venv"
 
-    conda create -n $CONDA_ENV python=3.10.0 -y
-    conda activate $CONDA_ENV
-    # This is optional if you prefer to use built-in nvcc
-    conda install -c nvidia cuda-toolkit -y
+# Check if the virtual environment needs to be created
+if [ ! -d "$VENV_DIR" ]; then
+    echo "Creating virtual environment in '$VENV_DIR' using uv..."
+    uv venv $VENV_DIR
 else
-    echo "Skipping conda environment creation. Make sure you have the correct environment activated."
+    echo "Virtual environment '$VENV_DIR' already exists. Skipping creation."
 fi
 
-# init a raw torch to avoid installation errors.
-# pip install torch
+# Activate the virtual environment
+echo "Activating virtual environment..."
+source $VENV_DIR/bin/activate
 
-# update pip to latest version for pyproject.toml setup.
-pip install -U pip
+# Update pip to the latest version
+echo "Updating pip to the latest version..."
+uv pip install -U pip
 
-# for fast attn
-pip install -U xformers==0.0.27.post2 --index-url https://download.pytorch.org/whl/cu121
-
-# install sana
-pip install -e .
+# Install required dependencies
+echo "Installing dependencies..."
+uv pip install -r pyproject.toml
 
 # install torchprofile
-# pip install git+https://github.com/zhijian-liu/torchprofile
+# uv pip install git+https://github.com/zhijian-liu/torchprofile
+
+echo "Environment setup completed."

pyproject.toml

@@ -1,80 +1,85 @@
 [build-system]
-requires = ["setuptools>=61.0"]
 build-backend = "setuptools.build_meta"
+requires = ["setuptools>=61.0"]
 
 [project]
-name = "sana"
-version = "0.0.1"
-description = "SANA"
-readme = "README.md"
-requires-python = ">=3.10"
 classifiers = [
-    "Programming Language :: Python :: 3",
-    "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: Apache Software License",
 ]
 dependencies = [
-    "pre-commit",
-    "accelerate",
-    "beautifulsoup4",
-    "bs4",
-    "came-pytorch",
-    "einops",
-    "ftfy",
-    "diffusers@git+https://github.com/huggingface/diffusers",
-    "clip@git+https://github.com/openai/CLIP.git",
-    "gradio",
-    "image-reward",
-    "ipdb",
-    "mmcv==1.7.2",
-    "omegaconf",
-    "opencv-python",
-    "optimum",
-    "patch_conv",
-    "peft",
-    "protobuf",
-    "pytorch-fid",
-    "regex",
-    "sentencepiece",
-    "tensorboard",
-    "tensorboardX",
-    "timm",
-    "torchaudio==2.4.0",
-    "torchvision==0.19",
-    "transformers",
-    "triton==3.0.0",
-    "wandb",
-    "webdataset",
-    "xformers==0.0.27.post2",
-    "yapf",
-    "spaces",
-    "matplotlib",
-    "termcolor",
-    "pyrallis",
+  "accelerate>=1.3.0",
+  "beautifulsoup4>=4.12.3",
+  "bs4>=0.0.2",
+  "came-pytorch>=0.1.3",
+  "clip@git+https://github.com/openai/CLIP.git",
+  "diffusers@git+https://github.com/huggingface/diffusers",
+  "einops>=0.8.0",
+  "ftfy>=6.3.1",
+  "gradio>=5.12.0",
+  "image-reward>=1.5",
+  "ipdb>=0.13.13",
+  "matplotlib>=3.10.0",
+  "mmcv==1.7.2",
+  "numpy>=1.26.4",
+  "omegaconf>=2.3.0",
+  "opencv-python>=4.11.0.86",
+  "optimum>=1.23.3",
+  "pandas>=2.2.3",
+  "patch-conv>=0.0.1b0",
+  "peft>=0.14.0",
+  "pre-commit>=4.0.1",
+  "protobuf>=5.29.3",
+  "pyrallis>=0.3.1",
+  "pytorch-fid>=0.3.0",
+  "regex>=2024.11.6",
+  "ruff>=0.9.2",
+  "sentencepiece>=0.2.0",
+  "spaces>=0.32.0",
+  "tensorboard>=2.18.0",
+  "tensorboardx>=2.6.2.2",
+  "termcolor>=2.5.0",
+  "timm>=0.6.13",
+  "torchaudio==2.4.0",
+  "torchvision==0.19",
+  "transformers>=4.48.0",
+  "triton==3.0.0",
+  "wandb>=0.19.4",
+  "webdataset>=0.2.100",
+  "xformers==0.0.27.post2",
+  "yapf>=0.43.0",
 ]
-
+description = "SANA"
+name = "sana"
+readme = "README.md"
+requires-python = ">=3.12"
+version = "0.0.1"
 
 [project.scripts]
 sana-run = "sana.cli.run:main"
 sana-upload = "sana.cli.upload2hf:main"
 
-[project.optional-dependencies]
-
-[project.urls]
+[tool.uv.sources]
+torch = {index = "pytorch-cu121"}
+torchaudio = {index = "pytorch-cu121"}
+torchvision = {index = "pytorch-cu121"}
 
-[tool.pip]
-extra-index-url = ["https://download.pytorch.org/whl/cu121"]
+[[tool.uv.index]]
+explicit = true
+name = "pytorch-cu121"
+url = "https://download.pytorch.org/whl/cu121"
 
-[tool.black]
-line-length = 120
+# [tool.black]
+# line-length = 120
 
-[tool.isort]
-profile = "black"
-multi_line_output = 3
-include_trailing_comma = true
-force_grid_wrap = 0
-use_parentheses = true
-ensure_newline_before_comments = true
-line_length = 120
+# [tool.isort]
+# profile = "black"
+# multi_line_output = 3
+# include_trailing_comma = true
+# force_grid_wrap = 0
+# use_parentheses = true
+# ensure_newline_before_comments = true
+# line_length = 120
 
 [tool.setuptools.packages.find]
 exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]

ruff.toml

@@ -0,0 +1,77 @@
+# Exclude a variety of commonly ignored directories.
+exclude = [
+  ".bzr",
+  ".direnv",
+  ".eggs",
+  ".git",
+  ".git-rewrite",
+  ".hg",
+  ".ipynb_checkpoints",
+  ".mypy_cache",
+  ".nox",
+  ".pants.d",
+  ".pyenv",
+  ".pytest_cache",
+  ".pytype",
+  ".ruff_cache",
+  ".svn",
+  ".tox",
+  ".venv",
+  ".vscode",
+  "__pypackages__",
+  "_build",
+  "buck-out",
+  "build",
+  "dist",
+  "node_modules",
+  "site-packages",
+  ".venv",
+]
+
+# Same as Black.
+indent-width = 4
+line-length = 120
+
+# Assume Python 3.12
+target-version = "py312"
+
+[lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+ignore = []
+select = ["E4", "E7", "E9", "F"]
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"