Refactor to deploy modal app.

Makefile

@@ -1,8 +1,11 @@
+deploy:
+	modal deploy sdcli.py
+
 run:
-	modal run sd_cli.py \
+	modal run entrypoint.py \
-	--prompt "A woman with bob hair" \
+	--prompt "a photograph of an astronaut riding a horse" \
 	--n-prompt "" \
-	--height 768 \
+	--height 512 \
 	--width 512 \
-	--samples 5 \
+	--samples 1 \
-	--steps 30
+	--steps 50

README.md

@@ -31,7 +31,8 @@ To use the script, execute the below.
 1. git clone the repository.
 2. Create the `.env` file and set a huggingface API token and a model with reference to `.env.example`.
 3. Open the Makefile and set prompts.
-4. Execute `make run` command.
+4. Execute `make deploy` command. An application will be deployed to Modal by the command.
+5. Execute `make run` command.
 
 Images are generated and output to the `outputs/` directory.
 

entrypoint.py

@@ -0,0 +1,56 @@
+import time
+
+import modal
+
+stub = modal.Stub("run-stable-diffusion-cli")
+stub.run_inference = modal.Function.from_name("stable-diffusion-cli", "StableDiffusion.run_inference")
+
+
+@stub.local_entrypoint()
+def main(
+    prompt: str,
+    n_prompt: str,
+    height: int = 512,
+    width: int = 512,
+    samples: int = 5,
+    batch_size: int = 1,
+    steps: int = 20,
+    seed: int = -1,
+):
+    """
+    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.
+    """
+    import util
+
+    directory = util.make_directory()
+    seed_generated = seed
+    for i in range(samples):
+        if seed == -1:
+            seed_generated = util.generate_seed()
+        start_time = time.time()
+        # images = sd.run_inference(seed=seed_generated)
+        images = stub.app.run_inference.call(
+            prompt=prompt,
+            n_prompt=n_prompt,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            steps=steps,
+            seed=seed_generated,
+        )
+        util.save_images(directory, images, seed_generated, i)
+        total_time = time.time() - start_time
+        print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
+
+    prompts: dict[str, int | str] = {
+        "prompt": prompt,
+        "n_prompt": n_prompt,
+        "height": height,
+        "width": width,
+        "samples": samples,
+        "batch_size": batch_size,
+        "steps": steps,
+    }
+    util.save_prompts(prompts)

sdcli.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import io
 import os
-import time
 from urllib.request import Request, urlopen
 
 from modal import Image, Mount, Secret, Stub, method
@@ -94,38 +93,20 @@ class StableDiffusion(ClsMixin):
     A class that wraps the Stable Diffusion pipeline and scheduler.
     """
 
-    def __init__(
+    def __enter__(self):
-        self,
-        prompt: str,
-        n_prompt: str,
-        height: int = 512,
-        width: int = 512,
-        samples: int = 1,
-        batch_size: int = 1,
-        steps: int = 30,
-    ):
         import diffusers
         import torch
 
-        self.prompt = prompt
-        self.n_prompt = n_prompt
-        self.height = height
-        self.width = width
-        self.samples = samples
-        self.batch_size = batch_size
-        self.steps = steps
         self.use_vae = os.environ["USE_VAE"] == "true"
         self.upscaler = os.environ["UPSCALER"]
         self.use_face_enhancer = os.environ["USE_FACE_ENHANCER"] == "true"
         self.use_hires_fix = os.environ["USE_HIRES_FIX"] == "true"
-
         self.cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
         if os.path.exists(self.cache_path):
             print(f"The directory '{self.cache_path}' exists.")
         else:
             print(f"The directory '{self.cache_path}' does not exist. Download models...")
             download_models()
-        self.max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt)
 
         torch.backends.cuda.matmul.allow_tf32 = True
 
@@ -203,23 +184,34 @@ class StableDiffusion(ClsMixin):
         return max_embeddings_multiples
 
     @method()
-    def run_inference(self, seed: int) -> list[bytes]:
+    def run_inference(
+        self,
+        prompt: str,
+        n_prompt: str,
+        height: int = 512,
+        width: int = 512,
+        samples: int = 1,
+        batch_size: int = 1,
+        steps: int = 30,
+        seed: int = 1,
+    ) -> list[bytes]:
         """
         Runs the Stable Diffusion pipeline on the given prompt and outputs images.
         """
         import torch
 
+        max_embeddings_multiples = self.count_token(p=prompt, n=n_prompt)
         generator = torch.Generator("cuda").manual_seed(seed)
         with torch.inference_mode():
             with torch.autocast("cuda"):
                 base_images = self.pipe.text2img(
-                    self.prompt * self.batch_size,
+                    prompt * batch_size,
-                    negative_prompt=self.n_prompt * self.batch_size,
+                    negative_prompt=n_prompt * batch_size,
-                    height=self.height,
+                    height=height,
-                    width=self.width,
+                    width=width,
-                    num_inference_steps=self.steps,
+                    num_inference_steps=steps,
                     guidance_scale=7.5,
-                    max_embeddings_multiples=self.max_embeddings_multiples,
+                    max_embeddings_multiples=max_embeddings_multiples,
                     generator=generator,
                 ).images
 
@@ -236,12 +228,12 @@ class StableDiffusion(ClsMixin):
                     with torch.inference_mode():
                         with torch.autocast("cuda"):
                             hires_fixed = self.pipe.img2img(
-                                prompt=self.prompt * self.batch_size,
+                                prompt=prompt * batch_size,
-                                negative_prompt=self.n_prompt * self.batch_size,
+                                negative_prompt=n_prompt * batch_size,
-                                num_inference_steps=self.steps,
+                                num_inference_steps=steps,
                                 strength=0.3,
                                 guidance_scale=7.5,
-                                max_embeddings_multiples=self.max_embeddings_multiples,
+                                max_embeddings_multiples=max_embeddings_multiples,
                                 generator=generator,
                                 image=img,
                             ).images
@@ -336,84 +328,3 @@ class StableDiffusion(ClsMixin):
         torch.cuda.empty_cache()
 
         return upscaled_imgs
-
-    # TODO: Implement this
-    # @method()
-    # def img2img(
-    #     self,
-    #     prompt: str,
-    #     n_prompt: str,
-    #     batch_size: int = 1,
-    #     steps: int = 20,
-    #     strength: float = 0.3,
-    #     max_embeddings_multiples: int = 1,
-    #     # image: Image.Image = None,
-    #     base_images: list[Image.Image],
-    # ):
-    #     import torch
-
-    #     torch.cuda.empty_cache()
-    #     for img in base_images:
-    #         with torch.inference_mode():
-    #             with torch.autocast("cuda"):
-    #                 hires_fixed = self.pipe.img2img(
-    #                     prompt=prompt * batch_size,
-    #                     negative_prompt=n_prompt * batch_size,
-    #                     num_inference_steps=steps],
-    #                     strength=strength,
-    #                     guidance_scale=7.5,
-    #                     max_embeddings_multiples=max_embeddings_multiples,
-    #                     generator=generator,
-    #                     image=img,
-    #                 ).images
-    #         base_images.extend(hires_fixed)
-    #     torch.cuda.empty_cache()
-
-
-@stub.local_entrypoint()
-def entrypoint(
-    prompt: str,
-    n_prompt: str,
-    height: int = 512,
-    width: int = 512,
-    samples: int = 5,
-    batch_size: int = 1,
-    steps: int = 20,
-    seed: int = -1,
-):
-    """
-    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.
-    """
-    import util
-
-    directory = util.make_directory()
-
-    sd = StableDiffusion.remote(
-        prompt=prompt,
-        n_prompt=n_prompt,
-        height=height,
-        width=width,
-        batch_size=batch_size,
-        steps=steps,
-    )
-    for i in range(samples):
-        if seed == -1:
-            seed_generated = util.generate_seed()
-        start_time = time.time()
-        images = sd.run_inference(seed=seed_generated)
-        util.save_images(directory, images, seed_generated, i)
-        total_time = time.time() - start_time
-        print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
-
-    prompts: dict[str, int | str] = {
-        "prompt": prompt,
-        "n_prompt": n_prompt,
-        "height": height,
-        "width": width,
-        "samples": samples,
-        "batch_size": batch_size,
-        "steps": steps,
-    }
-    util.save_prompts(prompts)