# https://blog.runpod.io/custom-models-with-serverless-and-automatic-stable-diffusion/
FROM runpod/stable-diffusion:web-automatic-base-4.0.0

SHELL ["/bin/bash", "-c"]

ENV PATH="${PATH}:/workspace/stable-diffusion-webui/venv/bin"

WORKDIR /workspace

# PFG: https://civitai.com/api/download/models/1316
# Hassanblend: https://civitai.com/api/download/models/4635
# Deliberate v2: https://civitai.com/api/download/models/15236

RUN wget -O /workspace/stable-diffusion-webui/models/Stable-diffusion/model.safetensors https://civitai.com/api/download/models/4635
## Extra downloads (for Hassanblend)
RUN wget -O /workspace/stable-diffusion-webui/models/VAE/vae-ft-mse-840000-ema-pruned.safetensors https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors

RUN mkdir /workspace/stable-diffusion-webui/models/Codeformer/ && \
    wget -O /workspace/stable-diffusion-webui/models/Codeformer/codeformer-v0.1.0.pth https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth && \
    wget -O /workspace/stable-diffusion-webui/repositories/CodeFormer/weights/facelib/detection_Resnet50_Final.pth https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth && \
    wget -O /workspace/stable-diffusion-webui/repositories/CodeFormer/weights/facelib/parsing_parsenet.pth https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth && \
    mkdir /workspace/stable-diffusion-webui/models/GFPGAN/ && \
    wget -O /workspace/stable-diffusion-webui/models/GFPGAN/GFPGANv1.4.pth https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth && \
    mkdir /workspace/stable-diffusion-webui/models/RealESRGAN/ && \
    wget -O /workspace/stable-diffusion-webui/models/RealESRGAN/RealESRGAN_x4plus.pth https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth && \
    mkdir /workspace/stable-diffusion-webui/models/ESRGAN/ && \
    wget -O /workspace/stable-diffusion-webui/models/ESRGAN/ESRGAN.pth https://github.com/cszn/KAIR/releases/download/v1.0/ESRGAN.pth && \
#    mkdir /workspace/stable-diffusion-webui/models/Stable-diffusion/tokenizer/ && \
#    wget -O /workspace/stable-diffusion-webui/models/Stable-diffusion/tokenizer/vocab.json https://huggingface.co/runwayml/stable-diffusion-v1-5/raw/main/tokenizer/vocab.json && \
    git clone --depth 1 https://github.com/Coyote-A/ultimate-upscale-for-automatic1111 /workspace/stable-diffusion-webui/extensions/ultimate-upscale-for-automatic1111
RUN pip install runpod

ADD handler.py /workspace/
ADD start.sh /workspace/
RUN chmod +x /workspace/start.sh

CMD [ "/workspace/start.sh" ]