losses fix

code/losses.py

@@ -14,12 +14,11 @@ from transformers import CLIPProcessor, CLIPModel
 from diffusers import StableDiffusionPipeline
 
 class SDSLoss(nn.Module):
-    def __init__(self, cfg, device):
+    def __init__(self, cfg, device, model):
         super(SDSLoss, self).__init__()
         self.cfg = cfg
         self.device = device
-        self.pipe = StableDiffusionPipeline.from_pretrained(cfg.diffusion.model,
-                                                       torch_dtype=torch.float16, use_auth_token=cfg.token)
+        self.pipe = model
         self.pipe = self.pipe.to(self.device)
 
         self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(self.device)
@@ -55,8 +54,6 @@ class SDSLoss(nn.Module):
             text_embeddings = img_emb
             uncond_embeddings = img_emb
 
-        print(text_embeddings.size())
-        print(uncond_embeddings.size())
         self.text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
         self.text_embeddings = self.text_embeddings.repeat_interleave(self.cfg.batch_size, 0)
         del self.pipe.tokenizer