Add application file4

house_diffusion/gaussian_diffusion.py

@@ -898,7 +898,8 @@ class GaussianDiffusion:
                 bin_target = bin_target * 256 #-> [0, 256]
                 bin_target = dec2bin(bin_target.permute([0,2,1]).round().int(), 8) 
                 bin_target = bin_target.reshape([target.shape[0], target.shape[2], 16]).permute([0,2,1])
-                t_weights = (t<10).cuda().unsqueeze(1).unsqueeze(2)
+                # t_weights = (t<10).cuda().unsqueeze(1).unsqueeze(2)
+                t_weights = (t<10).unsqueeze(1).unsqueeze(2)
                 t_weights = t_weights * (t_weights.shape[0]/max(1, t_weights.sum()))
                 bin_target[bin_target==0] = -1
                 assert model_output_bin.shape == bin_target.shape

house_diffusion/transformer.py

@@ -77,7 +77,7 @@ class MultiHeadAttention(nn.Module):
         q = q.transpose(1,2)
         v = v.transpose(1,2)# calculate attention using function we will define next
         #TODO
-        mask = mask.to('cuda:0')
+        # mask = mask.to('cuda:0')
         scores = attention(q, k, v, self.d_k, mask, self.dropout)
         # concatenate heads and put through final linear layer
         concat = scores.transpose(1,2).contiguous().view(bs, -1, self.d_model)
@@ -232,8 +232,8 @@ class TransformerModel(nn.Module):
 
         # Different input embeddings (Input, Time, Conditions)
         #TODO---------------------------------------------------------------
-        x = x.to('cuda:0')
-        timesteps = timesteps.to(x.device)
+        # x = x.to('cuda:0')
+        # timesteps = timesteps.to(x.device)
         # print(x.device)
 
         time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
@@ -247,7 +247,7 @@ class TransformerModel(nn.Module):
                 else:
                     cond = th.cat((cond, kwargs[key]), 2)
             #TODO
-            cond = cond.to('cuda:0')
+            # cond = cond.to('cuda:0')
             cond_emb = self.condition_emb(cond.float())
 
         # PositionalEncoding and DM model