Add sd_cli.py.

sd_cli.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+import io
+import os
+import time
+from datetime import date
+from pathlib import Path
+from modal import Image, Secret, Stub, method
+
+stub = Stub("stable-diffusion-cli")
+
+MODEL = {
+    "repo_id": "runwayml/stable-diffusion-v1-5",
+    "name": "stable-diffusion-v1-5",
+}
+CACHE_PATH = os.path.join("/vol/cache", MODEL["name"])
+
+
+def download_models():
+    """
+    Downloads the model from Hugging Face and saves it to the cache path using
+    diffusers.StableDiffusionPipeline.from_pretrained().
+    """
+    import diffusers
+    import torch
+
+    hugging_face_token = os.environ["HUGGINGFACE_TOKEN"]
+
+    scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
+        MODEL["repo_id"],
+        subfolder="scheduler",
+        use_auth_token=hugging_face_token,
+        cache_dir=CACHE_PATH,
+    )
+    scheduler.save_pretrained(CACHE_PATH, safe_serialization=True)
+
+    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_name("my-huggingface-secret")],
+    )
+)
+stub.image = stub_image
+
+
+@stub.cls(gpu="A10G", secrets=[Secret.from_name("my-huggingface-secret")])
+class StableDiffusion:
+    """
+    A class that wraps the Stable Diffusion pipeline and scheduler.
+    """
+
+    def __enter__(self):
+        import diffusers
+        import torch
+
+        if os.path.exists(CACHE_PATH):
+            print(f"The directory '{CACHE_PATH}' exists.")
+        else:
+            print(f"The directory '{CACHE_PATH}' does not exist. Download models...")
+            download_models()
+
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+        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",
+        ).to("cuda")
+        if self.pipe.safety_checker is not None:
+            self.pipe.safety_checker = lambda images, **kwargs: (images, False)
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+    @method()
+    def run_inference(
+        self,
+        prompt: str,
+        n_prompt: str,
+        steps: int = 30,
+        batch_size: int = 1,
+        height: int = 512,
+        width: int = 512,
+    ) -> list[bytes]:
+        """
+        Runs the Stable Diffusion pipeline on the given prompt and outputs images.
+        """
+        import torch
+
+        with torch.inference_mode():
+            with torch.autocast("cuda"):
+                images = self.pipe(
+                    [prompt] * batch_size,
+                    negative_prompt=[n_prompt] * batch_size,
+                    height=height,
+                    width=width,
+                    num_inference_steps=steps,
+                    guidance_scale=7.5,
+                ).images
+
+        # Convert to PNG bytes
+        image_output = []
+        for image in images:
+            with io.BytesIO() as buf:
+                image.save(buf, format="PNG")
+                image_output.append(buf.getvalue())
+        return image_output
+
+
+@stub.local_entrypoint()
+def entrypoint(
+    prompt: str,
+    n_prompt: str,
+    samples: int = 5,
+    steps: int = 30,
+    batch_size: int = 1,
+    height: int = 512,
+    width: int = 512,
+):
+    """
+    This function is the entrypoint for the Runway CLI.
+    The function pass the given prompt to StableDiffusion on Modal,
+    gets back a list of images and outputs images to local.
+
+    The function is called with the following arguments:
+    - prompt: the prompt to run inference on
+    - n_prompt: the negative prompt to run inference on
+    - samples: the number of samples to generate
+    - steps: the number of steps to run inference for
+    - batch_size: the batch size to use
+    - height: the height of the output image
+    - width: the width of the output image
+    """
+    print(f"steps => {steps}, sapmles => {samples}, batch_size => {batch_size}")
+
+    directory = Path(f"./outputs/{date.today().strftime('%Y-%m-%d')}")
+    if not directory.exists():
+        directory.mkdir(exist_ok=True, parents=True)
+
+    stable_diffusion = StableDiffusion()
+    for i in range(samples):
+        start_time = time.time()
+        images = stable_diffusion.run_inference.call(
+            prompt,
+            n_prompt,
+            steps,
+            batch_size,
+            height,
+            width,
+        )
+        total_time = time.time() - start_time
+        print(
+            f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)."
+        )
+        for j, image_bytes in enumerate(images):
+            output_path = directory / f"output_{j}_{i}.png"
+            print(f"Saving it to {output_path}")
+            with open(output_path, "wb") as file:
+                file.write(image_bytes)