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

Fix sample image gen to work with block swap

Browse Source
This commit is contained in:
Kohya S
2024-10-28 22:08:57 +09:00
parent 1065dd1b56
commit af8e216035
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
library/sd3_train_utils.py
View File

@@ -364,6 +364,7 @@ def do_sample(
x_c_nc = torch.cat([x, x], dim=0) x_c_nc = torch.cat([x, x], dim=0)
# print(x_c_nc.shape, timestep.shape, c_crossattn.shape, y.shape) # print(x_c_nc.shape, timestep.shape, c_crossattn.shape, y.shape)
mmdit.prepare_block_swap_before_forward()
model_output = mmdit(x_c_nc, timestep, context=c_crossattn, y=y) model_output = mmdit(x_c_nc, timestep, context=c_crossattn, y=y)
model_output = model_output.float() model_output = model_output.float()
batched = model_sampling.calculate_denoised(sigma_hat, model_output, x) batched = model_sampling.calculate_denoised(sigma_hat, model_output, x)
@@ -385,6 +386,7 @@ def do_sample(
x = x + d * dt x = x + d * dt
x = x.to(dtype) x = x.to(dtype)
mmdit.prepare_block_swap_before_forward()
return x return x