Update MagicQuill/llava_new.py

MagicQuill/llava_new.py

@@ -48,12 +48,19 @@ class LLaVAModel:
         image_sizes = []
         to_pil = transforms.ToPILImage()
         for image in images:
-            image = image.clone().permute(2, 0, 1).cpu()
-            image = to_pil(image)
-            image_sizes.append(image.size)
-            image = expand2square(image, tuple(int(x) for x in self.image_processor.image_mean))
-            image = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-            images_tensor.append(image.half())
+    image = image.clone().permute(2, 0, 1).cpu()
+    image = to_pil(image)
+    image_sizes.append(image.size)
+    image = expand2square(image, tuple(int(x) for x in self.image_processor.image_mean))
+    image = self.image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
+    
+    # ✅ Rzutowanie na float32 lub float16 zależnie od dostępności CUDA
+    if torch.cuda.is_available():
+        image = image.half()
+    else:
+        image = image.float()
+
+    images_tensor.append(image)
 
         conv = conv_templates["llava_v1"].copy()
         conv.append_message(conv.roles[0], qs)