From 337bf01048e7374b1c04f82adfaa1b4e296d84af Mon Sep 17 00:00:00 2001 From: hodanov <1031hoda@gmail.com> Date: Tue, 13 Jun 2023 00:10:48 +0900 Subject: [PATCH 1/2] Modify to use seed. --- Makefile | 3 ++- sd_cli.py | 12 +++++++++--- util.py | 17 ++++++++++++++--- 3 files changed, 25 insertions(+), 7 deletions(-) diff --git a/Makefile b/Makefile index 16c1533..a1c5ec5 100644 --- a/Makefile +++ b/Makefile @@ -2,8 +2,9 @@ run: modal run sd_cli.py \ --prompt "A woman with bob hair" \ --n-prompt "" \ + --upscaler "RealESRGAN_x4plus_anime_6B" \ --height 768 \ --width 512 \ --samples 5 \ --steps 50 \ - --upscaler "RealESRGAN_x4plus_anime_6B" + --seed 500 diff --git a/sd_cli.py b/sd_cli.py index edb70e8..005af60 100644 --- a/sd_cli.py +++ b/sd_cli.py @@ -101,6 +101,7 @@ class StableDiffusion: """ import torch + generator = torch.Generator("cuda").manual_seed(inputs["seed"]) with torch.inference_mode(): with torch.autocast("cuda"): base_images = self.pipe( @@ -111,6 +112,7 @@ class StableDiffusion: num_inference_steps=inputs["steps"], guidance_scale=7.5, max_embeddings_multiples=inputs["max_embeddings_multiples"], + generator=generator, ).images if inputs["upscaler"] != "": @@ -197,12 +199,13 @@ class StableDiffusion: def entrypoint( prompt: str, n_prompt: str, + upscaler: str, height: int = 512, width: int = 512, samples: int = 5, batch_size: int = 1, steps: int = 20, - upscaler: str = "", + seed: int = -1, ): """ This function is the entrypoint for the Runway CLI. @@ -210,6 +213,9 @@ def entrypoint( gets back a list of images and outputs images to local. """ + if seed == -1: + seed = util.generate_seed() + inputs: dict[str, int | str] = { "prompt": prompt, "n_prompt": n_prompt, @@ -219,7 +225,7 @@ def entrypoint( "batch_size": batch_size, "steps": steps, "upscaler": upscaler, - # seed=-1 + "seed": seed, } inputs["max_embeddings_multiples"] = util.count_token(p=prompt, n=n_prompt) @@ -229,7 +235,7 @@ def entrypoint( for i in range(samples): start_time = time.time() images = sd.run_inference.call(inputs) - util.save_images(directory, images, i) + util.save_images(directory, images, inputs, i) total_time = time.time() - start_time print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).") diff --git a/util.py b/util.py index affb2a0..40aa05f 100644 --- a/util.py +++ b/util.py @@ -1,4 +1,5 @@ """ Utility functions for the script. """ +import random import time from datetime import date from pathlib import Path @@ -9,6 +10,16 @@ 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. @@ -16,7 +27,7 @@ def make_directory() -> Path: directory = Path(f"{OUTPUT_DIRECTORY}/{DATE_TODAY}") if not directory.exists(): directory.mkdir(exist_ok=True, parents=True) - print(f"Make directory: {directory}") + print(f"Make a directory: {directory}") return directory @@ -54,13 +65,13 @@ def count_token(p: str, n: str) -> int: return max_embeddings_multiples -def save_images(directory: Path, images: list[bytes], i: int): +def save_images(directory: Path, images: list[bytes], inputs: dict[str, int | str], 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}_{i}_{j}.png" + output_path = directory / f"{formatted_time}_{inputs['seed']}_{i}_{j}.png" print(f"Saving it to {output_path}") with open(output_path, "wb") as file: file.write(image_bytes) From ff580d08e5ba9be804268ce87f4f4a63c8e171e5 Mon Sep 17 00:00:00 2001 From: hodanov <1031hoda@gmail.com> Date: Fri, 16 Jun 2023 19:58:04 +0900 Subject: [PATCH 2/2] Modify to add seed to filename. --- .env.example | 1 + Makefile | 3 +-- sd_cli.py | 39 +++++++++++++++++++++++---------------- util.py | 6 ++---- 4 files changed, 27 insertions(+), 22 deletions(-) diff --git a/.env.example b/.env.example index 98ef154..90a3abb 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,4 @@ HUGGING_FACE_TOKEN="" MODEL_REPO_ID="stabilityai/stable-diffusion-2-1" MODEL_NAME="stable-diffusion-2-1" +USE_VAE="false" diff --git a/Makefile b/Makefile index a1c5ec5..febe53e 100644 --- a/Makefile +++ b/Makefile @@ -6,5 +6,4 @@ run: --height 768 \ --width 512 \ --samples 5 \ - --steps 50 \ - --seed 500 + --steps 50 diff --git a/sd_cli.py b/sd_cli.py index 005af60..01b7a01 100644 --- a/sd_cli.py +++ b/sd_cli.py @@ -75,23 +75,31 @@ class StableDiffusion: torch.backends.cuda.matmul.allow_tf32 = True - vae = diffusers.AutoencoderKL.from_pretrained( - cache_path, - subfolder="vae", - ) - scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( cache_path, subfolder="scheduler", ) - self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( - cache_path, - scheduler=scheduler, - vae=vae, - custom_pipeline="lpw_stable_diffusion", - torch_dtype=torch.float16, - ).to("cuda") + if os.environ["USE_VAE"] == "true": + vae = diffusers.AutoencoderKL.from_pretrained( + cache_path, + subfolder="vae", + ) + self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( + cache_path, + scheduler=scheduler, + vae=vae, + custom_pipeline="lpw_stable_diffusion", + torch_dtype=torch.float16, + ).to("cuda") + else: + self.pipe = diffusers.StableDiffusionPipeline.from_pretrained( + cache_path, + scheduler=scheduler, + custom_pipeline="lpw_stable_diffusion", + torch_dtype=torch.float16, + ).to("cuda") + self.pipe.enable_xformers_memory_efficient_attention() @method() @@ -213,9 +221,6 @@ def entrypoint( gets back a list of images and outputs images to local. """ - if seed == -1: - seed = util.generate_seed() - inputs: dict[str, int | str] = { "prompt": prompt, "n_prompt": n_prompt, @@ -233,9 +238,11 @@ def entrypoint( sd = StableDiffusion() for i in range(samples): + if seed == -1: + inputs["seed"] = util.generate_seed() start_time = time.time() images = sd.run_inference.call(inputs) - util.save_images(directory, images, inputs, i) + util.save_images(directory, images, int(inputs["seed"]), i) total_time = time.time() - start_time print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).") diff --git a/util.py b/util.py index 40aa05f..c7c1d75 100644 --- a/util.py +++ b/util.py @@ -4,8 +4,6 @@ import time from datetime import date from pathlib import Path -from PIL import Image - OUTPUT_DIRECTORY = "outputs" DATE_TODAY = date.today().strftime("%Y-%m-%d") @@ -65,13 +63,13 @@ def count_token(p: str, n: str) -> int: return max_embeddings_multiples -def save_images(directory: Path, images: list[bytes], inputs: dict[str, int | str], i: int): +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}_{inputs['seed']}_{i}_{j}.png" + 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)