CPU mode: remove fp16 + autocast, use fp32 everywhere

code/losses.py

@@ -15,17 +15,26 @@ class SDSLoss(nn.Module):
         super(SDSLoss, self).__init__()
         self.cfg = cfg
         self.device = device
+        self.fp16 = device.type == "cuda"
+        dtype = torch.float16 if self.fp16 else torch.float32
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            cfg.diffusion.model,
+            torch_dtype=dtype,
+            token=cfg.token,
+        ).to(device)
+        
         self.pipe = StableDiffusionPipeline.from_pretrained(
             cfg.diffusion.model,
             torch_dtype=torch.float32,
             token=cfg.token,
         ).to("cpu")
 
-        #self.pipe.enable_xformers_memory_efficient_attention()
-        self.pipe.enable_attention_slicing(slice_size=1)
-        self.pipe.enable_vae_slicing()
-        self.pipe.enable_vae_tiling()
-        self.pipe.unet.enable_gradient_checkpointing()
+        if self.fp16:
+            # self.pipe.enable_xformers_memory_efficient_attention()
+            self.pipe.enable_attention_slicing(slice_size=1)
+            self.pipe.enable_vae_slicing()
+            self.pipe.enable_vae_tiling()
+            self.pipe.unet.enable_gradient_checkpointing()
 
         alphas_cumprod = torch.tensor(self.pipe.scheduler.alphas_cumprod)
         self.alphas = alphas_cumprod.to(device)
@@ -38,21 +47,20 @@ class SDSLoss(nn.Module):
         #self.pipe.enable_model_cpu_offload()
 
         # text-encoder is no longer needed
-        del self.pipe.text_encoder, self.pipe.tokenizer
+        #del self.pipe.text_encoder, self.pipe.tokenizer
 
     def embed_text(self):
         tok = self.pipe.tokenizer
-        txt   = tok(self.cfg.caption, padding="max_length",
-                    max_length=tok.model_max_length,
-                    truncation=True, return_tensors="pt")
-        un    = tok([""], padding="max_length",
-                    max_length=tok.model_max_length,
-                    return_tensors="pt")
+        txt = tok(self.cfg.caption, padding="max_length",
+                  max_length=tok.model_max_length, truncation=True,
+                  return_tensors="pt")
+        un  = tok([""], padding="max_length",
+                  max_length=tok.model_max_length, return_tensors="pt")
 
         with torch.no_grad():
-            te = self.pipe.text_encoder.eval()          # still real tensors
-            em_txt = te(txt.input_ids   ).last_hidden_state.to(torch.float16)
-            em_un  = te(un .input_ids   ).last_hidden_state.to(torch.float16)
+            te = self.pipe.text_encoder.eval()
+            em_txt = te(txt.input_ids).last_hidden_state.to(torch.float32)
+            em_un  = te(un .input_ids).last_hidden_state.to(torch.float32)
 
         self.text_embeddings = (
             torch.cat([em_un, em_txt])
@@ -64,10 +72,14 @@ class SDSLoss(nn.Module):
 
     def forward(self, x_aug: torch.Tensor) -> torch.Tensor:
         # ---------------------------------------------------- encode
-        x = (x_aug * 2.0 - 1.0).to(self.device, dtype=torch.float16)
-        with torch.cuda.amp.autocast():
+        x = (x_aug * 2.0 - 1.0).to(self.device, dtype=torch.float32)
+        if self.fp16:
+            with torch.cuda.amp.autocast():
+                latents = self.pipe.vae.encode(x).latent_dist.sample()
+        else:
             latents = self.pipe.vae.encode(x).latent_dist.sample()
-        latents = 0.18215 * latents.to(self.device, dtype=torch.float16)
+
+        latents = 0.18215 * latents
         torch.cuda.empty_cache() 
 
         # ---------------------------------------------------- add noise