diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..464b462 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,5 @@ +FROM python:3.11.3-slim-bullseye +COPY requirements.txt / +RUN apt update \ + && apt install -y wget git \ + && pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117 --pre xformers diff --git a/requirements.txt b/requirements.txt index 0f5b34d..98c8d5e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,9 @@ -modal-client -# The below packages are not necessary because them are installed into the image on Modal. -# However, you prevent some lint errors when you edit the code by the editor like VSCode or Vim. -diffusers -torch +accelerate +scipy +diffusers[torch] +safetensors +torch==2.0.1+cu117 +torchvision +torchmetrics +omegaconf +transformers diff --git a/sd_cli.py b/sd_cli.py index 7c67ea9..70e3e52 100644 --- a/sd_cli.py +++ b/sd_cli.py @@ -4,7 +4,7 @@ import os import time from datetime import date from pathlib import Path -from modal import Image, Secret, Stub, method +from modal import Image, Secret, Stub, method, Mount stub = Stub("stable-diffusion-cli") @@ -17,7 +17,6 @@ def download_models(): diffusers.StableDiffusionPipeline.from_pretrained(). """ import diffusers - import torch hugging_face_token = os.environ["HUGGINGFACE_TOKEN"] model_repo_id = os.environ["MODEL_REPO_ID"] @@ -34,30 +33,17 @@ def download_models(): pipe = diffusers.StableDiffusionPipeline.from_pretrained( model_repo_id, use_auth_token=hugging_face_token, - torch_dtype=torch.float16, cache_dir=cache_path, ) pipe.save_pretrained(cache_path, safe_serialization=True) -stub_image = ( - Image.debian_slim(python_version="3.10") - .pip_install( - "accelerate", - "diffusers[torch]>=0.15.1", - "ftfy", - "torch", - "torchvision", - "transformers~=4.25.1", - "triton", - "safetensors", - "torch>=2.0", - ) - .pip_install("xformers", pre=True) - .run_function( - download_models, - secrets=[Secret.from_dotenv(__file__)], - ) +stub_image = Image.from_dockerfile( + path="./Dockerfile", + context_mount=Mount.from_local_file("./requirements.txt"), +).run_function( + download_models, + secrets=[Secret.from_dotenv(__file__)], ) stub.image = stub_image @@ -84,21 +70,14 @@ class StableDiffusion: scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( cache_path, subfolder="scheduler", - solver_order=2, - prediction_type="epsilon", - thresholding=False, - algorithm_type="dpmsolver++", - solver_type="midpoint", - denoise_final=True, # important if steps are <= 10 - low_cpu_mem_usage=True, - device_map="auto", ) self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( cache_path, scheduler=scheduler, - low_cpu_mem_usage=True, - device_map="auto", + custom_pipeline="lpw_stable_diffusion", + max_embeddings_multiples=2, + safety_checker=None, ).to("cuda") self.pipe.enable_xformers_memory_efficient_attention()