Update model.py

model.py

@@ -1121,15 +1121,24 @@ class StreamMultiDiffusion(nn.Module):
         else:
             x_t_latent_plus_uc = x_t_latent  # (T * p, 4, h, w)
 
-        try:
-            model_pred = self.unet(
-                x_t_latent_plus_uc.to(self.unet.dtype),  # (B, 4, h, w)
-                t_list,  # (B,)
-                encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
-                return_dict=False,
-            )[0]  # (B, 4, h, w)
-        except Exception as e:
-            print(e)
+        ns = []
+        c1, c2, c3 = 0, 0, 0
+        for n, p in self.unet.named_parameters():
+            if p.data.dtype == torch.float:
+                c1 += 1
+                ns.append(n)
+            elif p.data.dtype == torch.half:
+                c2 += 1
+            else:
+                c3 += 1
+        print(c1, c2, c3)
+        print(ns)
+        model_pred = self.unet(
+            x_t_latent_plus_uc.to(self.unet.dtype),  # (B, 4, h, w)
+            t_list,  # (B,)
+            encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
+            return_dict=False,
+        )[0]  # (B, 4, h, w)
         print('222222222222222', model_pred.dtype)
 
         if self.bootstrap_steps[0] > 0: