FROM tensorflow/tensorflow:devel-gpu
ENV DEBIAN_FRONTEND=noninteractive
ENV NVIDIA_VISIBLE_DEVICES=all
ENV CUDA_VISIBLE_DEVICES=0
ENV PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.6,max_split_size_mb:128
RUN apt update && apt install -y build-essential cmake git aria2 python3 cuda-runtime-11-8 && apt clean
RUN git clone --depth 1 https://github.com/basujindal/stable-diffusion /stable-diffusion
RUN pip install --upgrade pip
RUN pip install transformers==4.19.2 diffusers invisible-watermark omegaconf einops torchvision pytorch_lightning clip kornia
WORKDIR /stable-diffusion
RUN pip install -e .
RUN git clone --depth 1 https://github.com/CompVis/taming-transformers && mv taming-transformers/taming . && rm -rf taming-transformers
RUN aria2c --seed-time=0 "magnet:?xt=urn:btih:2daef5b5f63a16a9af9169a529b1a773fc452637&dn=v1-5-pruned-emaonly.ckpt&tr=udp%3a%2f%2ftracker.opentrackr.org%3a1337%2fannounce&tr=udp%3a%2f%2f9.rarbg.com%3a2810%2fannounce&tr=udp%3a%2f%2ftracker.openbittorrent.com%3a6969%2fannounce&tr=udp%3a%2f%2fopentracker.i2p.rocks%3a6969%2fannounce&tr=https%3a%2f%2fopentracker.i2p.rocks%3a443%2fannounce&tr=http%3a%2f%2ftracker.openbittorrent.com%3a80%2fannounce&tr=udp%3a%2f%2ftracker.torrent.eu.org%3a451%2fannounce&tr=udp%3a%2f%2fopen.stealth.si%3a80%2fannounce&tr=udp%3a%2f%2fvibe.sleepyinternetfun.xyz%3a1738%2fannounce&tr=udp%3a%2f%2ftracker2.dler.org%3a80%2fannounce&tr=udp%3a%2f%2ftracker1.bt.moack.co.kr%3a80%2fannounce&tr=udp%3a%2f%2ftracker.zemoj.com%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.tiny-vps.com%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.theoks.net%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.publictracker.xyz%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.monitorit4.me%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.moeking.me%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.lelux.fi%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.dler.org%3a6969%2fannounce&tr=udp%3a%2f%2ftracker.army%3a6969%2fannounce"
RUN mkdir -p models/ldm/stable-diffusion-v1/ && ln -s $PWD/v1-5-pruned-emaonly.ckpt $PWD/models/ldm/stable-diffusion-v1/model.ckpt
RUN sed -i "s/import argparse, os, re/import argparse, os, sys, re\nsys.path.insert(0, os.getcwd())/" optimizedSD/optimized_txt2img.py optimizedSD/optimized_img2img.py
RUN python optimizedSD/optimized_txt2img.py 2> /dev/null || exit 0
ENTRYPOINT ["python", "optimizedSD/optimized_txt2img.py", "--H", "512", "--W", "512", "--n_samples", "1", "--precision", "full", "--outdir", "/output", "--format", "jpg", "--prompt"]