Update train.py

train.py

@@ -31,56 +31,43 @@ def load_model(model_name, device_id=0):
 
 def caption_batch(batch, processor, model):
     images = batch["image"]
-
-    processed_images = []
+    
+    pil_images = []
     for image in images:
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
         if image.mode != "RGB":
             image = image.convert("RGB")
-        processed_images.append(image)
+        pil_images.append(image)
 
-    encoded_list = []
-    for image in processed_images:
+    text_inputs = []
+    for _ in pil_images:
         msg = [
             {
                 "role": "user",
                 "content": [
-                    {"type": "image", "image": image},
-                    {
-                        "type": "text",
-                        "text": "Describe the image, and skip mentioning that it's illustrated or from anime.",
-                    },
+                    {"type": "text", "text": "Describe the image, and skip mentioning that it's illustrated or from anime."},
                 ],
             }
         ]
+        text_inputs.append(processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True))
 
-        enc = processor.apply_chat_template(
-            msg,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_dict=True,
-            return_tensors="pt",
-        )
-
-        encoded_list.append(enc)
-
-    input_ids = torch.nn.utils.rnn.pad_sequence(
-        [e.input_ids[0] for e in encoded_list],
-        batch_first=True,
-        padding_value=processor.tokenizer.pad_token_id,
-    ).to(model.device)
+    inputs = processor(
+        text=text_inputs,
+        images=pil_images,
+        return_tensors="pt",
+        padding=True
+    )
 
-    attention_mask = torch.nn.utils.rnn.pad_sequence(
-        [e.attention_mask[0] for e in encoded_list],
-        batch_first=True,
-        padding_value=0,
-    ).to(model.device)
+    input_ids = inputs.input_ids.to(model.device)
+    attention_mask = inputs.attention_mask.to(model.device)
+    pixel_values = inputs.pixel_values.to(model.device)
 
     with torch.no_grad():
         generated = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
+            pixel_values=pixel_values,
             max_new_tokens=256,
         )
 
@@ -90,16 +77,16 @@ def caption_batch(batch, processor, model):
     for d in decoded:
         if "<|im_start|>assistant" in d:
             d = d.split("<|im_start|>assistant")[-1].strip()
-
+        
         special_tokens = set(processor.tokenizer.all_special_tokens)
         for token in special_tokens:
             d = d.replace(token, "")
-
+        
         d = d.strip()
         captions.append(d)
 
     return {
-        "image": processed_images,
+        "image": images,
         "text": captions,
     }