This commit is contained in:
lvmin
2023-08-09 15:07:09 -07:00
parent c56e2c9fd1
commit c104f30137
3 changed files with 4 additions and 2 deletions
+1 -1
View File
@@ -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] * 255.0
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)
cv2.imwrite('img.png', samples)