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

Refactoring.

Browse Source
This commit is contained in:
hodanov 2023-06-17 18:11:45 +09:00
parent a5adb6b9fb
commit b0862045ca
1 changed files with 10 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
sd_cli.py
View File

@ -104,30 +104,24 @@ class StableDiffusion:
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True
scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained( self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
cache_path,
custom_pipeline="lpw_stable_diffusion",
torch_dtype=torch.float16,
)
self.pipe.scheduler = diffusers.EulerAncestralDiscreteScheduler.from_pretrained(
cache_path, cache_path,
subfolder="scheduler", subfolder="scheduler",
) )
if os.environ["USE_VAE"] == "true": if os.environ["USE_VAE"] == "true":
vae = diffusers.AutoencoderKL.from_pretrained( self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
cache_path, cache_path,
subfolder="vae", subfolder="vae",
) )
self.pipe = diffusers.StableDiffusionPipeline.from_pretrained(
cache_path, self.pipe.to("cuda")
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")
if os.environ["LORA_NAMES"] != "": if os.environ["LORA_NAMES"] != "":
lora_names = os.getenv("LORA_NAMES").split(",") lora_names = os.getenv("LORA_NAMES").split(",")