Fix aesthetic score (again), add llava reward

ddpo_pytorch/aesthetic_scorer.py

@@ -30,22 +30,22 @@ class MLP(nn.Module):
 
 
 class AestheticScorer(torch.nn.Module):
-    def __init__(self):
+    def __init__(self, dtype):
         super().__init__()
         self.clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
         self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
         self.mlp = MLP()
         state_dict = torch.load(ASSETS_PATH.joinpath("sac+logos+ava1-l14-linearMSE.pth"))
         self.mlp.load_state_dict(state_dict)
+        self.dtype = dtype
         self.eval()
 
     @torch.no_grad()
     def __call__(self, images):
-        assert isinstance(images, list)
-        assert isinstance(images[0], Image.Image)
+        device = next(self.parameters()).device
         inputs = self.processor(images=images, return_tensors="pt")
-        inputs = {k: v.cuda() for k, v in inputs.items()}
+        inputs = {k: v.to(self.dtype).to(device) for k, v in inputs.items()}
         embed = self.clip.get_image_features(**inputs)
         # normalize embedding
         embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)
-        return self.mlp(embed)
+        return self.mlp(embed).squeeze(1)