Issue fix

app.py

@@ -357,22 +357,38 @@ class SkinGPT4(nn.Module):
         #     "### Response:"
         # )
 
-        prompt_parts = [
-            "### Instruction: <Img>",
-            "<Image>",
-            "</Img> Could you describe the skin disease in this image for me? ### Response:"
-        ]
+        # prompt_parts = [
+        #     "### Instruction: <Img>",
+        #     "<Image>",
+        #     "</Img> Could you describe the skin disease in this image for me? ### Response:"
+        # ]
+        prompt = "### Instruction: <Img><Image></Img> Could you describe the skin disease in this image for me? ### Response:"
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
+
         # Tokenize each part separately
-        tokens_before = self.tokenizer(prompt_parts[0], return_tensors="pt").input_ids.to(images.device)
-        tokens_after = self.tokenizer(prompt_parts[2], return_tensors="pt").input_ids.to(images.device)
-        input_ids = torch.cat([
-            tokens_before[:, :-1],  # Remove EOS from first part
-            torch.full((1, 1), self.tokenizer.convert_tokens_to_ids("<Image>")).to(images.device),
-            tokens_after[:, 1:]  # Remove BOS from second part
-        ], dim=1)
-        embeddings = self.llama.model.embed_tokens(input_ids)
-        image_token_pos = (input_ids == self.tokenizer.convert_tokens_to_ids("<Image>")).nonzero()
-        embeddings[image_token_pos] = aligned_features.mean(dim=1)
+        # tokens_before = self.tokenizer(prompt_parts[0], return_tensors="pt").input_ids.to(images.device)
+        # tokens_after = self.tokenizer(prompt_parts[2], return_tensors="pt").input_ids.to(images.device)
+        # input_ids = torch.cat([
+        #     tokens_before[:, :-1],  # Remove EOS from first part
+        #     torch.full((1, 1), self.tokenizer.convert_tokens_to_ids("<Image>")).to(images.device),
+        #     tokens_after[:, 1:]  # Remove BOS from second part
+        # ], dim=1)
+        # embeddings = self.llama.model.embed_tokens(input_ids)
+        # image_token_pos = (input_ids == self.tokenizer.convert_tokens_to_ids("<Image>")).nonzero()
+        # embeddings[image_token_pos] = aligned_features.mean(dim=1)
+
+        image_token_id = self.tokenizer.convert_tokens_to_ids("<Image>")
+        image_token_pos = (inputs.input_ids == image_token_id).nonzero()
+
+        if image_token_pos.shape[0] != 1:
+            raise ValueError(f"Expected 1 image token, found {image_token_pos.shape[0]}")
+
+        # Prepare embeddings
+        embeddings = self.llama.model.embed_tokens(inputs.input_ids)
+
+        row, col = image_token_pos[0]
+        embeddings[row, col] = aligned_features.mean(dim=1)
+
         outputs = self.llama.generate(
             inputs_embeds=embeddings,
             max_length=max_length,