Modify a directory structure.

sdcli/txt2img.py

@@ -0,0 +1,61 @@
+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,
+    upscaler: str = "",
+    use_face_enhancer: str = "False",
+    use_hires_fix: str = "False",
+):
+    """
+    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 = stub.app.run_inference.call(
+            prompt=prompt,
+            n_prompt=n_prompt,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            steps=steps,
+            seed=seed_generated,
+            upscaler=upscaler,
+            use_face_enhancer=use_face_enhancer == "True",
+            use_hires_fix=use_hires_fix == "True",
+        )
+        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/util.py

@@ -0,0 +1,55 @@
+""" Utility functions for the script. """
+import random
+import time
+from datetime import date
+from pathlib import Path
+
+OUTPUT_DIRECTORY = "outputs"
+DATE_TODAY = date.today().strftime("%Y-%m-%d")
+
+
+def generate_seed() -> int:
+    """
+    Generate a random seed.
+    """
+    seed = random.randint(0, 4294967295)
+    print(f"Generate a random seed: {seed}")
+
+    return seed
+
+
+def make_directory() -> Path:
+    """
+    Make a directory for saving outputs.
+    """
+    directory = Path(f"{OUTPUT_DIRECTORY}/{DATE_TODAY}")
+    if not directory.exists():
+        directory.mkdir(exist_ok=True, parents=True)
+        print(f"Make a directory: {directory}")
+
+    return directory
+
+
+def save_prompts(inputs: dict):
+    """
+    Save prompts to a file.
+    """
+    prompts_filename = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+    with open(
+        file=f"{OUTPUT_DIRECTORY}/{DATE_TODAY}/prompts_{prompts_filename}.txt", mode="w", encoding="utf-8"
+    ) as file:
+        for name, value in inputs.items():
+            file.write(f"{name} = {repr(value)}\n")
+        print(f"Save prompts: {prompts_filename}.txt")
+
+
+def save_images(directory: Path, images: list[bytes], seed: int, i: int):
+    """
+    Save images to a file.
+    """
+    for j, image_bytes in enumerate(images):
+        formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+        output_path = directory / f"{formatted_time}_{seed}_{i}_{j}.png"
+        print(f"Saving it to {output_path}")
+        with open(output_path, "wb") as file:
+            file.write(image_bytes)