lllyasviel commited on
Commit
fbbf2aa
·
1 Parent(s): 60fe69a
Files changed (2) hide show
  1. entry.py +1 -1
  2. sgm/modules/diffusionmodules/util.py +3 -1
entry.py CHANGED
@@ -151,6 +151,6 @@ with torch.no_grad():
151
  model.first_stage_model.cpu()
152
 
153
  import cv2
154
- samples = einops.rearrange(samples, 'b c h w -> b h w c')[0] * 255.0
155
  samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)
156
  cv2.imwrite('img.png', samples)
 
151
  model.first_stage_model.cpu()
152
 
153
  import cv2
154
+ samples = einops.rearrange(samples, 'b c h w -> b h w c')[0, :, :, ::-1] * 255.0
155
  samples = samples.cpu().numpy().clip(0, 255).astype(np.uint8)
156
  cv2.imwrite('img.png', samples)
sgm/modules/diffusionmodules/util.py CHANGED
@@ -272,7 +272,9 @@ class SiLU(nn.Module):
272
 
273
  class GroupNorm32(nn.GroupNorm):
274
  def forward(self, x):
275
- return super().forward(x)
 
 
276
 
277
 
278
  def conv_nd(dims, *args, **kwargs):
 
272
 
273
  class GroupNorm32(nn.GroupNorm):
274
  def forward(self, x):
275
+ self.weight.to(torch.float32)
276
+ self.bias.to(torch.float32)
277
+ return super().forward(x.float()).type(x.dtype)
278
 
279
 
280
  def conv_nd(dims, *args, **kwargs):