chunzhi/stable-diffusion-cli-on-modal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix to switch Stable Diffusion depending on versions.

Browse Source
This commit is contained in:
hodanov 2024-11-04 08:23:14 +09:00
parent a4a161ad83
commit df6caf8f55
1 changed files with 44 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
cmd/sdxl_txt2img.py
View File

@ -1,18 +1,42 @@
from __future__ import annotations
import logging import logging
import time import time
import domain
import modal import modal
from domain import OutputDirectory, Prompts, Seed, StableDiffusionOutputManger
app = modal.App("run-stable-diffusion-cli") from infrasctucture import RunInferenceInterface, RunInferenceSDXLTxt2Img
run_inference = modal.Function.from_name(
"stable-diffusion-cli",
"SDXLTxt2Img.run_inference",
)
@app.local_entrypoint() def new_run_inference(
version: str,
prompts: Prompts,
output_format: str,
*,
use_upscaler: bool,
) -> RunInferenceInterface:
match version:
case "sd15":
# TODO: sd15用のクラスを実装したら置き換える
return RunInferenceSDXLTxt2Img(
prompts=prompts,
use_upscaler=use_upscaler,
output_format=output_format,
)
case "sdxl":
return RunInferenceSDXLTxt2Img(
prompts=prompts,
use_upscaler=use_upscaler,
output_format=output_format,
)
case _:
msg = f"Invalid version: {version}. Must be 'sd15' or 'sdxl'."
raise ValueError(msg)
@modal.App("run-stable-diffusion-cli").local_entrypoint()
def main( def main(
version: str,
prompt: str, prompt: str,
n_prompt: str, n_prompt: str,
height: int = 1024, height: int = 1024,
@ -34,27 +58,24 @@ def main(
) )
logger = logging.getLogger("run-stable-diffusion-cli") logger = logging.getLogger("run-stable-diffusion-cli")
output_directory = domain.OutputDirectory() output_directory = OutputDirectory()
directory_path = output_directory.make_directory() directory_path = output_directory.make_directory()
logger.info("Made a directory: %s", directory_path) logger.info("Made a directory: %s", directory_path)
prompts = domain.Prompts(prompt, n_prompt, height, width, samples, steps) prompts = Prompts(prompt, n_prompt, height, width, samples, steps)
sd_output_manager = domain.StableDiffusionOutputManger(prompts, directory_path) sd_output_manager = StableDiffusionOutputManger(prompts, directory_path)
run_inference = new_run_inference(
version,
prompts,
output_format,
use_upscaler=use_upscaler == "True",
)
for sample_index in range(samples): for sample_index in range(samples):
new_seed = domain.Seed(seed)
start_time = time.time() start_time = time.time()
images = run_inference.remote( new_seed = Seed(seed)
prompt=prompt, images = run_inference.exec(new_seed)
n_prompt=n_prompt,
height=height,
width=width,
steps=steps,
seed=new_seed.value,
use_upscaler=use_upscaler == "True",
output_format=output_format,
)
for generated_image_index, image_bytes in enumerate(images): for generated_image_index, image_bytes in enumerate(images):
saved_path = sd_output_manager.save_image( saved_path = sd_output_manager.save_image(
image_bytes, image_bytes,
@ -64,7 +85,6 @@ def main(
output_format, output_format,
) )
logger.info("Saved image to the: %s", saved_path) logger.info("Saved image to the: %s", saved_path)
total_time = time.time() - start_time total_time = time.time() - start_time
logger.info("Sample %s, took %ss (%ss / image).", sample_index, total_time, (total_time) / len(images)) logger.info("Sample %s, took %ss (%ss / image).", sample_index, total_time, (total_time) / len(images))