https://github.com/dataiku/dss-plugin-ai-art
DSS plugin to generate images from text using Stable Diffusion
https://github.com/dataiku/dss-plugin-ai-art
Last synced: 6 days ago
JSON representation
DSS plugin to generate images from text using Stable Diffusion
- Host: GitHub
- URL: https://github.com/dataiku/dss-plugin-ai-art
- Owner: dataiku
- License: apache-2.0
- Created: 2022-12-15T08:56:39.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2023-10-07T22:33:32.000Z (over 2 years ago)
- Last Synced: 2023-11-08T16:45:30.245Z (over 2 years ago)
- Language: Python
- Homepage:
- Size: 1.91 MB
- Stars: 1
- Watchers: 19
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
Awesome Lists containing this project
README
# AI Art
AI Art is a plugin for Dataiku that allows you to generate images from text prompts.
## Usage
With this plugin, you will be able to:
- Generate images from a prompt using **Stable Diffusion**
- Generate images from another image and a prompt using **Stable Diffusion**
## How to set up
### CUDA
A CUDA-capable GPU with at least 6 GB of VRAM is recommended.
You can run the plugin without a GPU, but it will be significantly slower.
If you’re using a CUDA GPU, your GPU must support CUDA 11.6, and the NVIDIA drivers must be installed on the server.
For installation instructions, see [NVIDIA Driver Installation Quickstart Guide](https://docs.nvidia.com/datacenter/tesla/tesla-installation-notes/index.html).
### Download weights
Before you can use the plugin, you need to download pre-trained weights to a local managed folder. The plugin is designed to work with version 2 of the Stable Diffusion weights, although it may work with other versions and derivatives as well.
For example, if you run the below code in a notebook, it will download the Stability AI version 2.1 weights, and upload them to a new managed folder.
The code requires a code environment with the *huggingface-hub* package installed.
##### Limitations and warnings
- The weights available from Stability AI are licensed under the
CreativeML OpenRAIL++-M license, which restricts usage. You can view the
license [here](https://huggingface.co/stabilityai/stable-diffusion-2/raw/main/LICENSE-MODEL). Hence, the weights must be acquired manually before you can use this plugin.
- The weights must be stored on the local filesystem. If a remote folder (S3,
etc) is used, or if the recipe uses containerized execution, the weights
will be downloaded to a temporary directory every time the recipe is run.
This is because the method used to load the weights ([from_pretrained])
requires a local filepath.
[from_pretrained]: https://huggingface.co/docs/diffusers/v0.6.0/en/api/diffusion_pipeline#diffusers.DiffusionPipeline.from_pretrained
```
import tempfile
import dataiku
import huggingface_hub
# Name of the local folder that will be created
FOLDER_NAME = "my_weights"
# Hugging Face repository
REPO = "stabilityai/stable-diffusion-2-1"
# Revision of the repository
#
# If using the Stability AI weights, the "fp16" revision is recommended
# if using CUDA. Otherwise, use the "main" revision.
REVISION = "fp16"
client = dataiku.api_client()
project = client.get_default_project()
# Create the managed folder
folder = project.create_managed_folder(FOLDER_NAME)
with tempfile.TemporaryDirectory(dir=".") as temp_dir:
# Download the weights to a temp local dir
huggingface_hub.snapshot_download(
repo_id=REPO,
revision=REVISION,
local_dir=temp_dir,
local_dir_use_symlinks=False,
# Skip large files that the plugin doesn't need
ignore_patterns=["*.ckpt", "*.safetensors"],
)
# Upload the weights to the managed folder
folder.upload_folder("/", temp_dir)
```
Using a folder that’s stored on the local filesystem is recommended. If the folder is stored on a remote connection (Amazon S3, Google Cloud Storage, etc), the weights will be downloaded to a temporary directory every time the recipe is run.
## How to use
AI Art contains two methods for generating images: Text-to-Image
Generation, and Text-Guided Image-to-Image Generation
### Text-to-Image Generation
Text-to-Image Generation is used to generate images from a text prompt.
1. Create a *Text-to-Image Generation* recipe with your weights folderas the input.
2. Enter a text prompt in the *Prompt* field.
3. If you’re using the *fp16* revision of the weights, be sure to check the *Half precision* field.
### Text-Guided Image-to-Image Generation
Text-Guided Image-to-Image Generation is used to modify an existing
reference image based on a text prompt.
1. Obtain a reference image that you want to use as a base, and upload it to a managed folder.
2. Create a *Text-Guided Image-to-Image Generation* recipe with your weights folder and your base-image folder as the inputs.
3. Enter a text prompt in the *Prompt* field.
4. Enter the path to your base image in the *Base image* field.
5. If you’re using the *fp16* revision of the weights, be sure to check the *Half precision* field.
## Testing and linting
Format the code using Black:
```bash
make black
```
Lint the code using Flake8, and verify that the code was formatted using Black:
```bash
make lint
```
Run unit tests:
```bash
make unit-tests
```
Run integration tests:
```bash
make integration-tests
```
Run all linters and tests at once:
```bash
make tests
```
## Building
Create a plugin archive that can be imported into Dataiku:
```bash
make plugin
```
## Known limitations
1. Due to licensing restrictions, the weights must be acquired manually before
you can use this plugin. See the [documentation][plugin_documentation] for
details.
1. The weights must be stored on the local filesystem. If a remote folder (S3,
etc) is used, or if the recipe uses containerized execution, the weights
will be downloaded to a temporary directory every time the recipe is run.
This is because the method used to load the weights ([from_pretrained])
requires a local filepath.
[from_pretrained]: https://huggingface.co/docs/diffusers/v0.6.0/en/api/diffusion_pipeline#diffusers.DiffusionPipeline.from_pretrained