disco-diffusion-wrapper

一个 disco diffusion 包装器的实现，可在您自己的 GPU 上运行批量文本输入。

基于这项工作，我们构建了一个 AI 绘画网站（https://6pen.art/），您可以尝试使用。

本仓库的主要工作

1. 将初始代码的模型加载和模型推理部分分离。现在您可以像这样使用：

# init and load pretrain model.
disco = DiscoDiffusion()
# infer
disco.draw(text, ...)

2. 使用 deepL 对文本进行预处理，这样你就可以使用任何喜欢的语言进行绘制。 默认语言对为 Chinese->English（ZH->EN-US），你可以在 run.py 或 run_batch.py 中进行更改。

def translate(text, source_lang="ZH", target_lang="EN-US"):
    res = translator.translate_text(text, source_lang="ZH", target_lang="EN-US")
    return res

您还可以通过设置 USE_TRANSLATE=False 来关闭文本翻译功能，这样就不需要 DeepL 的 authKey 了。

3. 批量生成与保存。 由于加载部分和推理部分是分离的，因此对于新句子，您无需再次加载预训练模型。这节省了大量时间。

cd disco-diffusion-wrapper/wrapper
# create and fill the sentence.list
python run_batch.py

输出图像将以原始文本命名保存，因此即使经过多次操作后，您也能轻松找到所需图像。

使用方法

git clone git@github.com:mazzzystar/disco-diffusion-wrapper.git
cd disco-diffusion-wrapper/
conda create --name disco
conda activate disco
pip install -r requirements.txt
cd wrapper/

# download pretrain model checkpoints
"""If this step is slow, you can download these model mannually from urls in the 
code of mutils.py, then put these checkpoints in the corresponding folder."""
python mutils.py

# Get your deepL API key, then add it to your run.py or batch_run.py to replace:
USE_TRANSLATE=True
auth_key = ""

# generated image by one sentence
python run.py "一行白鹭上青天"

# batch generation
touch sentence.list
# fill your text in, one by a line.
python run_batch.py

示例

以下是使用本仓库生成的部分示例。

半亩方塘一鉴开，天光云影共徘徊，水墨画

雪山，平原，幽鬼军队在夜里打着荧光绿色的火把前行，抽象画，Trending on artstation

躺在废墟中的玻璃鲸鱼，抽象画

明月松间照，清泉石上流，中国画

Elf Castle

山随平野尽，江入大荒流，Trending on artstation

您可以从这里下载我们所有的中国诗词生成结果。

如果您使用本项目并产生了一些有趣的成果，欢迎投稿。

问题

由于本项目仅为娱乐目的，我并未深入研究细节，并且删除了许多功能（如VR/3D/视频等），以便更快、更清晰地进行项目重构。如果您有兴趣恢复原始功能，非常欢迎提交PR。