This commit is contained in:
lvmin
2023-08-09 15:02:24 -07:00
parent ab65a357ec
commit 9993090813
+26 -27
View File
@@ -98,60 +98,59 @@ model.load_state_dict(safetensors.torch.load_file('./sd_xl_base_1.0.safetensors'
# model.conditioner.cuda() # model.conditioner.cuda()
model.conditioner.embedders[0].device = 'cpu' with torch.no_grad():
model.conditioner.embedders[1].device = 'cpu'
value_dict = { model.conditioner.embedders[0].device = 'cpu'
model.conditioner.embedders[1].device = 'cpu'
value_dict = {
"prompt": "a handsome man in forest", "negative_prompt": "ugly, bad", "orig_height": 1024, "orig_width": 1024, "prompt": "a handsome man in forest", "negative_prompt": "ugly, bad", "orig_height": 1024, "orig_width": 1024,
"crop_coords_top": 0, "crop_coords_left": 0, "target_height": 1024, "target_width": 1024, "aesthetic_score": 7.5, "crop_coords_top": 0, "crop_coords_left": 0, "target_height": 1024, "target_width": 1024, "aesthetic_score": 7.5,
"negative_aesthetic_score": 2.0, "negative_aesthetic_score": 2.0,
} }
batch, batch_uc = get_batch( batch, batch_uc = get_batch(
get_unique_embedder_keys_from_conditioner(model.conditioner), get_unique_embedder_keys_from_conditioner(model.conditioner),
value_dict, value_dict,
[1], [1],
) )
c, uc = model.conditioner.get_unconditional_conditioning( c, uc = model.conditioner.get_unconditional_conditioning(
batch, batch,
batch_uc=batch_uc) batch_uc=batch_uc)
# model.conditioner.cpu() # model.conditioner.cpu()
c = {a: b.to(torch.float16) for a, b in c.items()} c = {a: b.to(torch.float16) for a, b in c.items()}
uc = {a: b.to(torch.float16) for a, b in uc.items()} uc = {a: b.to(torch.float16) for a, b in uc.items()}
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
shape = (1, 4, 128, 128) shape = (1, 4, 128, 128)
randn = torch.randn(shape).to(torch.float16).cuda() randn = torch.randn(shape).to(torch.float16).cuda()
def denoiser(input, sigma, c): def denoiser(input, sigma, c):
return model.denoiser(model.model, input, sigma, c) return model.denoiser(model.model, input, sigma, c)
with torch.no_grad(): with torch.no_grad():
model.model.to(torch.float16).cuda() model.model.to(torch.float16).cuda()
model.denoiser.to(torch.float16).cuda() model.denoiser.to(torch.float16).cuda()
samples_z = sampler(denoiser, randn, cond=c, uc=uc) samples_z = sampler(denoiser, randn, cond=c, uc=uc)
model.model.cpu() model.model.cpu()
model.denoiser.cpu() model.denoiser.cpu()
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
a = 0 with torch.no_grad():
model.first_stage_model.to(torch.float16).cuda()
with torch.no_grad():
model.first_stage_model.cuda()
samples_x = model.decode_first_stage(samples_z) samples_x = model.decode_first_stage(samples_z)
samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0) samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
model.first_stage_model.cpu() model.first_stage_model.cpu()
import cv2 import cv2
samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 127.5 + 127.5 samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 127.5 + 127.5
samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8) samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)
cv2.imwrite('img.png', samples) cv2.imwrite('img.png', samples)