73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
import logging
|
|
import time
|
|
|
|
import domain
|
|
import modal
|
|
|
|
app = modal.App("run-stable-diffusion-cli")
|
|
run_inference = modal.Function.from_name(
|
|
"stable-diffusion-cli",
|
|
"SD15.run_txt2img_inference",
|
|
)
|
|
|
|
|
|
@app.local_entrypoint()
|
|
def main(
|
|
prompt: str,
|
|
n_prompt: str,
|
|
height: int = 512,
|
|
width: int = 512,
|
|
samples: int = 5,
|
|
steps: int = 20,
|
|
seed: int = -1,
|
|
use_upscaler: str = "",
|
|
fix_by_controlnet_tile: str = "False",
|
|
output_format: str = "png",
|
|
) -> None:
|
|
"""main() is the entrypoint for the Runway CLI.
|
|
This pass the given prompt to StableDiffusion on Modal, gets back a list of images and outputs images to local.
|
|
"""
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="[%(levelname)s] %(asctime)s - %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
logger = logging.getLogger("run-stable-diffusion-cli")
|
|
|
|
output_directory = domain.OutputDirectory()
|
|
directory_path = output_directory.make_directory()
|
|
logger.info("Made a directory: %s", directory_path)
|
|
|
|
prompts = domain.Prompts(prompt, n_prompt, height, width, samples, steps)
|
|
sd_output_manager = domain.StableDiffusionOutputManger(prompts, directory_path)
|
|
for sample_index in range(samples):
|
|
new_seed = domain.Seed(seed)
|
|
start_time = time.time()
|
|
images = run_inference.remote(
|
|
prompt=prompt,
|
|
n_prompt=n_prompt,
|
|
height=height,
|
|
width=width,
|
|
batch_size=1,
|
|
steps=steps,
|
|
seed=new_seed.value,
|
|
use_upscaler=use_upscaler == "True",
|
|
fix_by_controlnet_tile=fix_by_controlnet_tile == "True",
|
|
output_format=output_format,
|
|
)
|
|
for generated_image_index, image_bytes in enumerate(images):
|
|
saved_path = sd_output_manager.save_image(
|
|
image_bytes,
|
|
new_seed.value,
|
|
sample_index,
|
|
generated_image_index,
|
|
output_format,
|
|
)
|
|
logger.info("Saved image to the: %s", saved_path)
|
|
|
|
total_time = time.time() - start_time
|
|
logger.info("Sample %s, took %ss (%ss / image).", sample_index, total_time, (total_time) / len(images))
|
|
|
|
saved_prompts_path = sd_output_manager.save_prompts()
|
|
logger.info("Saved prompts: %s", saved_prompts_path)
|