Skip to content

Commit 8cec0d3

Browse files
ZachNagengastZachNagengast
authored
Merge pull request #4 from argmaxinc/swift-scaffolding
2 parents 9520836 + 223c717 commit 8cec0d3

33 files changed

+437
-68
lines changed

‎.flake8

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
[flake8]
2+
max-line-length = 120
3+
filename = *.py
4+
5+
[isort]
6+
profile = black

‎.github/workflows/lint.yml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
name: Check Lint
2+
3+
on:
4+
push:
5+
branches: [main]
6+
pull_request:
7+
branches: [main]
8+
9+
jobs:
10+
lint:
11+
runs-on: ubuntu-latest
12+
13+
steps:
14+
- name: Checkout code
15+
uses: actions/checkout@v3
16+
17+
- name: Set up Homebrew
18+
id: set-up-homebrew
19+
uses: Homebrew/actions/setup-homebrew@master
20+
21+
- name: Setup environment
22+
run: |
23+
make setup
24+
25+
- name: Run pre-commit hooks
26+
run: make format

‎.gitignore

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -231,4 +231,4 @@ fastlane/test_output
231231
!*.xcodeproj/project.pbxproj
232232
!*.xcodeproj/xcshareddata/
233233
!*.xcworkspace/contents.xcworkspacedata
234-
/*.gcno
234+
/*.gcno

‎.pre-commit-config.yaml

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
repos:
2+
- repo: https://github.com/pycqa/isort
3+
rev: 5.13.2
4+
hooks:
5+
- id: isort
6+
args: ["--profile", "black"]
7+
8+
- repo: https://github.com/psf/black
9+
rev: 23.3.0
10+
hooks:
11+
- id: black
12+
name: black
13+
language: python
14+
15+
- repo: https://github.com/pre-commit/pre-commit-hooks
16+
rev: v4.5.0
17+
hooks:
18+
- id: end-of-file-fixer
19+
- id: trailing-whitespace
20+
21+
- repo: local
22+
hooks:
23+
- id: swift-format
24+
name: swift-format
25+
entry: swift-format format --in-place --recursive .
26+
language: system
27+
types: [swift]
28+
files: \.swift$

‎.swift-format

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
{
2+
"version": 1,
3+
"indentation": {
4+
"spaces": 4
5+
},
6+
"spacesAroundRangeFormationOperators": true
7+
}

‎CODE_OF_CONDUCT.md

Whitespace-only changes.

‎CONTRIBUTING.md

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# Contributing to DiffusionKit
2+
3+
## Overview
4+
5+
We welcome and encourage contributions to DiffusionKit! Whether you're fixing bugs, improving documentation, or adding new features from the roadmap, your help is appreciated. This guide will help you get started with contributing to DiffusionKit.
6+
7+
## Getting Started
8+
9+
1. **Fork the Repository**: Start by [forking](https://github.com/argmaxinc/DiffusionKit/fork) the DiffusionKit repository on GitHub to your personal account.
10+
11+
2. **Clone Your Fork**: Clone your fork to your local machine to start making changes.
12+
13+
```bash
14+
git clone https://github.com/[your-username]/DiffusionKit.git
15+
cd DiffusionKit
16+
```
17+
18+
## Setting Up Your Development Environment
19+
20+
1. **Install Dependencies**: Use the provided `Makefile` to set up your environment. Run `make setup` to install necessary dependencies.
21+
22+
```bash
23+
make setup
24+
```
25+
26+
This will add any necessary CLI tools, as well as the pre-commit hooks for code formatting.
27+
28+
2. **Download Models**: Run to download the required models to run and test locally.
29+
30+
You can specify the model version as follows:
31+
32+
For python and mlx:
33+
```bash
34+
make download-model MODEL=stabilityai/stable-diffusion-3-medium
35+
```
36+
37+
For Swift with Core ML:
38+
```bash
39+
make download-model MODEL=argmaxinc/coreml-stable-diffusion-3-medium
40+
```
41+
42+
43+
## Making Changes
44+
45+
1. **Create a Branch**: Create a new branch for your changes.
46+
47+
```bash
48+
git checkout -b my-descriptive-branch-name
49+
```
50+
51+
2. **Make Your Changes**: Implement your changes, add new features, or fix bugs. Ensure you adhere to the existing coding style. If you're adding new features, make sure to update or add any documentation or tests as needed.
52+
53+
## Submitting Your Changes
54+
55+
1. **Commit Your Changes**: Once you're satisfied with your changes, commit them with a clear and concise commit message.
56+
57+
```bash
58+
git commit -am "Add a new feature"
59+
```
60+
61+
2. **Push to Your Fork**: Push your changes to your fork on GitHub.
62+
63+
```bash
64+
git push origin my-descriptive-branch-name
65+
```
66+
67+
3. **Create a Pull Request**: Go to the DiffusionKit repository on GitHub and create a new pull request from your fork. Ensure your pull request has a clear title and description.
68+
69+
4. **Code Review**: Wait for the maintainers to review your pull request. Be responsive to feedback and make any necessary changes.
70+
71+
## Guidelines
72+
73+
- **Code Style**: Follow the existing code style in the project.
74+
- **Commit Messages**: Write meaningful commit messages that clearly describe the changes.
75+
- **Documentation**: Update documentation if you're adding new features or making changes that affect how users interact with DiffusionKit.
76+
- **Tests**: Add or update tests for new features or bug fixes.
77+
78+
## Final Steps
79+
80+
After your pull request has been reviewed and approved, a maintainer will merge it into the main branch. Congratulations, you've successfully contributed to DiffusionKit!
81+
82+
Thank you for making DiffusionKit better for everyone! ❤️‍🔥

‎Makefile

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
.PHONY: setup setup-model-repo download-model format clean-package-caches
2+
3+
PIP_COMMAND := pip3
4+
PYTHON_COMMAND := python3
5+
6+
# Define model repository and directories
7+
MODEL_REPO_DIR := ./models/
8+
9+
setup:
10+
@echo "Setting up environment..."
11+
@which $(PIP_COMMAND)
12+
@which $(PYTHON_COMMAND)
13+
@echo "Checking for Homebrew..."
14+
@which brew > /dev/null || (echo "Error: Homebrew is not installed. Install it form here https://brew.sh and try again" && exit 1)
15+
@echo "Homebrew is installed."
16+
@echo "Checking for huggingface-cli..."
17+
@which huggingface-cli > /dev/null || (echo "Installing huggingface-cli..." && brew install huggingface-cli)
18+
@echo "huggingface-cli is installed."
19+
@echo "Checking for git-lfs..."
20+
@which git-lfs > /dev/null || (echo "Installing git-lfs..." && brew install git-lfs)
21+
@echo "git-lfs is installed."
22+
@if [ "$$(uname)" = "Darwin" ]; then \
23+
which trash > /dev/null || (echo "Installing trash..." && brew install trash); \
24+
echo "trash is installed."; \
25+
else \
26+
echo "Skipping trash installation (not macOS)."; \
27+
fi
28+
@echo "Checking for swift-format..."
29+
@which swift-format > /dev/null || (echo "Installing swift-format..." && brew install swift-format)
30+
@echo "swift-format is installed."
31+
@echo "Checking for pre-commit..."
32+
@which pre-commit > /dev/null || (echo "Installing pre-commit..." && brew install pre-commit && pre-commit install)
33+
@echo "pre-commit is installed."
34+
@echo "Done 🚀"
35+
36+
37+
# Download a specific model
38+
download-model:
39+
@if [ -z "$(MODEL)" ]; then \
40+
echo "Error: MODEL is not set. Usage: make download-model MODEL=your-model-repo"; \
41+
exit 1; \
42+
fi
43+
@echo "Setting up repository..."
44+
@mkdir -p $(MODEL_REPO_DIR)
45+
@if [ -d "$(MODEL_REPO_DIR)/$(MODEL)/.git" ]; then \
46+
echo "Repository exists, resetting..."; \
47+
export GIT_LFS_SKIP_SMUDGE=1; \
48+
cd $(MODEL_REPO_DIR)/$(MODEL) && git fetch --all && git reset --hard origin/main && git clean -fdx; \
49+
else \
50+
echo "Repository not found, initializing..."; \
51+
export GIT_LFS_SKIP_SMUDGE=1; \
52+
git clone https://huggingface.co/$(MODEL) $(MODEL_REPO_DIR)/$(MODEL); \
53+
fi
54+
@echo "Downloading model $(MODEL)..."
55+
@cd $(MODEL_REPO_DIR)/$(MODEL) && \
56+
git lfs pull --include="*"
57+
58+
format:
59+
@pre-commit run --all-files
60+
61+
clean-package-caches:
62+
@trash ~/Library/Caches/org.swift.swiftpm/repositories
63+
@trash ~/Library/Developer/Xcode/DerivedData

‎Package.resolved

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
{
2+
"originHash" : "138c786ec831c5ad1635043170e453f3f3d9e61e5f596ee9f733d3e9953a28d1",
3+
"pins" : [
4+
{
5+
"identity" : "mlx-swift",
6+
"kind" : "remoteSourceControl",
7+
"location" : "https://github.com/ml-explore/mlx-swift",
8+
"state" : {
9+
"branch" : "main",
10+
"revision" : "36d63a1fc386a551df14f5b67df1756dc17d2ebc"
11+
}
12+
},
13+
{
14+
"identity" : "swift-argument-parser",
15+
"kind" : "remoteSourceControl",
16+
"location" : "https://github.com/apple/swift-argument-parser.git",
17+
"state" : {
18+
"revision" : "c8ed701b513cf5177118a175d85fbbbcd707ab41",
19+
"version" : "1.3.0"
20+
}
21+
},
22+
{
23+
"identity" : "swift-numerics",
24+
"kind" : "remoteSourceControl",
25+
"location" : "https://github.com/apple/swift-numerics",
26+
"state" : {
27+
"revision" : "0a5bc04095a675662cf24757cc0640aa2204253b",
28+
"version" : "1.0.2"
29+
}
30+
},
31+
{
32+
"identity" : "swift-transformers",
33+
"kind" : "remoteSourceControl",
34+
"location" : "https://github.com/huggingface/swift-transformers.git",
35+
"state" : {
36+
"revision" : "fc6543263e4caed9bf6107466d625cfae9357f08",
37+
"version" : "0.1.8"
38+
}
39+
}
40+
],
41+
"version" : 3
42+
}

‎Package.swift

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
// swift-tools-version: 5.10
2+
// The swift-tools-version declares the minimum version of Swift required to build this package.
3+
4+
import PackageDescription
5+
6+
let package = Package(
7+
name: "DiffusionKit",
8+
platforms: [
9+
.macOS("13.3"),
10+
.iOS(.v16),
11+
.visionOS(.v1),
12+
],
13+
products: [
14+
.library(
15+
name: "DiffusionKit",
16+
targets: ["DiffusionKit"]
17+
),
18+
.library(
19+
name: "DiffusionKitMLX",
20+
targets: ["DiffusionKitMLX"]
21+
),
22+
],
23+
dependencies: [
24+
.package(url: "https://github.com/huggingface/swift-transformers.git", exact: "0.1.8"),
25+
.package(url: "https://github.com/apple/swift-argument-parser.git", exact: "1.3.0"),
26+
.package(url: "https://github.com/ml-explore/mlx-swift", branch: "main"),
27+
],
28+
targets: [
29+
.target(
30+
name: "DiffusionKit",
31+
dependencies: [
32+
.product(name: "Transformers", package: "swift-transformers")
33+
],
34+
path: "swift/Sources/DiffusionKit"
35+
),
36+
.target(
37+
name: "DiffusionKitMLX",
38+
dependencies: [
39+
.product(name: "MLX", package: "mlx-swift"),
40+
.product(name: "MLXFFT", package: "mlx-swift"),
41+
.product(name: "MLXNN", package: "mlx-swift"),
42+
.product(name: "Transformers", package: "swift-transformers"),
43+
],
44+
path: "swift/Sources/DiffusionKitMLX"
45+
),
46+
.testTarget(
47+
name: "DiffusionKitTests",
48+
dependencies: [
49+
"DiffusionKit",
50+
.product(name: "Transformers", package: "swift-transformers"),
51+
],
52+
path: "swift/Tests/DiffusionKitTests"
53+
),
54+
.testTarget(
55+
name: "DiffusionKitMLXTests",
56+
dependencies: [
57+
"DiffusionKitMLX",
58+
.product(name: "Transformers", package: "swift-transformers"),
59+
],
60+
path: "swift/Tests/DiffusionKitMLXTests"
61+
),
62+
]
63+
)

‎README.md

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,22 @@ Please refer to the help menu for all available arguments: `diffusionkit-cli -h`
9292
🚧
9393

9494
</details>
95+
96+
## License
97+
98+
DiffusionKit is released under the MIT License. See [LICENSE](LICENSE) for more details.
99+
100+
## Citation
101+
102+
If you use DiffusionKit for something cool or just find it useful, please drop us a note at [info@takeargmax.com](mailto:info@takeargmax.com)!
103+
104+
If you use DiffusionKit for academic work, here is the BibTeX:
105+
106+
```bibtex
107+
@misc{diffusionkit-argmax,
108+
title = {DiffusionKit},
109+
author = {Argmax, Inc.},
110+
year = {2024},
111+
URL = {https://github.com/argmaxinc/DiffusionKit}
112+
}
113+
```

‎python/src/mlx/__init__.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,29 +5,29 @@
55
# Copyright (C) 2024 Argmax, Inc. All Rights Reserved.
66
#
77

8-
import time
8+
import gc
99
import math
10+
import time
1011
from typing import Optional, Tuple
11-
import gc
1212

1313
import mlx.core as mx
1414
import mlx.nn as nn
1515
import numpy as np
16-
from PIL import Image
1716
from argmaxtools.utils import get_logger
17+
from PIL import Image
18+
1819
from python.src.utils import bytes2gigabytes
1920

2021
from .model_io import (
2122
_DEFAULT_MODEL,
23+
load_mmdit,
24+
load_t5_encoder,
25+
load_t5_tokenizer,
2226
load_text_encoder,
2327
load_tokenizer,
24-
load_mmdit,
2528
load_vae_decoder,
2629
load_vae_encoder,
27-
load_t5_encoder,
28-
load_t5_tokenizer,
2930
)
30-
3131
from .sampler import ModelSamplingDiscreteFlow
3232

3333
logger = get_logger(__name__)

0 commit comments

Comments
 (0)