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

Add token_count. Modify file name about generated images.

Browse Source
This commit is contained in:
hodanov 2023-05-24 20:47:33 +09:00
parent 51f2836df8
commit e313420b28
1 changed files with 14 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
sd_cli.py
View File

@ -76,7 +76,6 @@ class StableDiffusion:
cache_path, cache_path,
scheduler=scheduler, scheduler=scheduler,
custom_pipeline="lpw_stable_diffusion", custom_pipeline="lpw_stable_diffusion",
max_embeddings_multiples=2,
safety_checker=None, safety_checker=None,
).to("cuda") ).to("cuda")
self.pipe.enable_xformers_memory_efficient_attention() self.pipe.enable_xformers_memory_efficient_attention()
@ -90,6 +89,7 @@ class StableDiffusion:
batch_size: int = 1, batch_size: int = 1,
height: int = 512, height: int = 512,
width: int = 512, width: int = 512,
max_embeddings_multiples: int = 1,
) -> list[bytes]: ) -> list[bytes]:
""" """
Runs the Stable Diffusion pipeline on the given prompt and outputs images. Runs the Stable Diffusion pipeline on the given prompt and outputs images.
@ -105,6 +105,7 @@ class StableDiffusion:
width=width, width=width,
num_inference_steps=steps, num_inference_steps=steps,
guidance_scale=7.5, guidance_scale=7.5,
max_embeddings_multiples=max_embeddings_multiples,
).images ).images
# Convert to PNG bytes # Convert to PNG bytes
@ -142,6 +143,15 @@ def entrypoint(
""" """
print(f"steps => {steps}, sapmles => {samples}, batch_size => {batch_size}") print(f"steps => {steps}, sapmles => {samples}, batch_size => {batch_size}")
max_embeddings_multiples = 1
token_count = len(prompt.split())
if token_count > 77:
max_embeddings_multiples = token_count // 77 + 1
print(
f"token_count => {token_count}, max_embeddings_multiples => {max_embeddings_multiples}"
)
directory = Path(f"./outputs/{date.today().strftime('%Y-%m-%d')}") directory = Path(f"./outputs/{date.today().strftime('%Y-%m-%d')}")
if not directory.exists(): if not directory.exists():
directory.mkdir(exist_ok=True, parents=True) directory.mkdir(exist_ok=True, parents=True)
@ -156,13 +166,15 @@ def entrypoint(
batch_size, batch_size,
height, height,
width, width,
max_embeddings_multiples,
) )
total_time = time.time() - start_time total_time = time.time() - start_time
print( print(
f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)." f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image)."
) )
for j, image_bytes in enumerate(images): for j, image_bytes in enumerate(images):
output_path = directory / f"output_{j}_{i}.png" formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
output_path = directory / f"{formatted_time}_{i}_{j}.png"
print(f"Saving it to {output_path}") print(f"Saving it to {output_path}")
with open(output_path, "wb") as file: with open(output_path, "wb") as file:
file.write(image_bytes) file.write(image_bytes)