From 9cdf7671e2fa9e37ba983210c48a05333d8ea7a1 Mon Sep 17 00:00:00 2001 From: lvmin Date: Wed, 9 Aug 2023 15:13:08 -0700 Subject: [PATCH] i --- entry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/entry.py b/entry.py index 5eff7031..4f573e9c 100644 --- a/entry.py +++ b/entry.py @@ -151,6 +151,6 @@ with torch.no_grad(): model.first_stage_model.cpu() import cv2 - samples = einops.rearrange(samples, 'b c h w -> b h w c')[0, :, :, ::-1] * 255.0 - samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8) + samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 255.0 + samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)[:, :, ::-1] cv2.imwrite('img.png', samples)