VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-06 13:47:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix latent upscale not working if bs>1

Browse Source
This commit is contained in:
Kohya S
2023-04-25 08:10:21 +09:00
parent 1890535d1b
commit a85fcfe05f
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
gen_img_diffusers.py
View File

@@ -945,7 +945,7 @@ class PipelineLike:
# encode the init image into latents and scale the latents
init_image = init_image.to(device=self.device, dtype=latents_dtype)
if init_image.size()[1:] == (height // 8, width // 8):
if init_image.size()[-2:] == (height // 8, width // 8):
init_latents = init_image
else:
if vae_batch_size >= batch_size:

8
tools/latent_upscaler.py
View File

@@ -243,7 +243,13 @@ def create_upscaler(**kwargs):
model = Upscaler()
print(f"Loading weights from {weights}...")
model.load_state_dict(torch.load(weights, map_location=torch.device("cpu")))
if os.path.splitext(weights)[1] == ".safetensors":
from safetensors.torch import load_file
sd = load_file(weights)
else:
sd = torch.load(weights, map_location=torch.device("cpu"))
model.load_state_dict(sd)
return model