fix added

SkinGPT.py

@@ -225,14 +225,36 @@ class SkinGPT4(nn.Module):
     def generate(self, images, user_input=None, max_new_tokens=300):
 
         image_embeds = self.encode_image(images)
+
         print(f"Aligned features : {image_embeds}")
+        print(f"\n Images embeddings shape : {image_embeds.shape} \n Llama config hidden size : {self.llama.config.hidden_size}")
+        if image_embeds.shape[-1] != self.llama.config.hidden_size:
+            raise ValueError(
+                f"Feature dimension mismatch. "
+                f"Q-Former output: {image_embeds.shape[-1]}, "
+                f"LLaMA expected: {self.llama.config.hidden_size}"
+            )
 
-        prompt = (
-            "### Instruction: <Img><ImageHere></Img> "
-            "Could you describe the skin condition in this image? "
-            "Please provide a detailed analysis including possible diagnoses. "
-            "### Response:"
-        )
+
+        # prompt = (
+        #     "### Instruction: <Img><ImageHere></Img> "
+        #     "Could you describe the skin condition in this image? "
+        #     "Please provide a detailed analysis including possible diagnoses. "
+        #     "### Response:"
+        # )
+
+        prompt = """### Skin Diagnosis Protocol ###
+        <IMAGE>
+        Patient Presentation: [Describe visible symptoms]
+        Primary Differential Diagnosis:
+        1. 
+        2. 
+        3. 
+        Recommended Diagnostic Tests:
+        - 
+        Treatment Options:
+        - 
+        <|endoftext|>"""
 
         print(f"\n[DEBUG] Raw Prompt:\n{prompt}")
 
@@ -241,7 +263,8 @@ class SkinGPT4(nn.Module):
             token=token,
             padding_side="right"
         )
-        self.tokenizer.add_special_tokens({'additional_special_tokens': ['<Img>', '</Img>', '<ImageHere>']})
+        # self.tokenizer.add_special_tokens({'additional_special_tokens': ['<Img>', '</Img>', '<ImageHere>']})
+        self.tokenizer.add_tokens(["<IMAGE>"])
         self.llama.resize_token_embeddings(len(self.tokenizer))
 
         inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
@@ -249,17 +272,24 @@ class SkinGPT4(nn.Module):
         print(f"\n[DEBUG] Tokenized input IDs:\n{inputs.input_ids}")
         print(f"[DEBUG] Special token positions: {self.tokenizer.all_special_tokens}")
 
-        image_token_id = self.tokenizer.convert_tokens_to_ids("<ImageHere>")
+        # image_token_id = self.tokenizer.convert_tokens_to_ids("<ImageHere>")
+        image_token_id = self.tokenizer.convert_tokens_to_ids("<IMAGE>")
+
+
+        # Prepare embeddings
+        input_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
+        visual_embeds = image_embeds.mean(dim=1, keepdim=True)
+
+        visual_embeds = F.layer_norm(visual_embeds, [visual_embeds.size(-1)])
         image_token_pos = torch.where(inputs.input_ids == image_token_id)
+
         if len(image_token_pos[0]) == 0:
             raise ValueError("Image token not found in prompt")
 
         print(f"\n[DEBUG] Image token found at position: {image_token_pos}")
 
-
-        # Prepare embeddings
-        input_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
-        visual_embeds = image_embeds.mean(dim=1, keepdim=True)
+        for pos in zip(*image_token_pos):
+            input_embeddings[pos] = visual_embeds[0, 0]
 
         print(f"\n[DEBUG] Before replacement:")
         print(f"Text embeddings shape: {input_embeddings.shape}")
@@ -269,7 +299,7 @@ class SkinGPT4(nn.Module):
 
         if visual_embeds.dtype != input_embeddings.dtype:
             visual_embeds = visual_embeds.to(input_embeddings.dtype)
-        input_embeddings[image_token_pos] = visual_embeds
+        # input_embeddings[image_token_pos] = visual_embeds
 
         print(f"\n[DEBUG] After replacement:")
         print(f"Image token embedding (after):\n{input_embeddings[0, image_token_pos[1], :5]}...")
@@ -309,11 +339,13 @@ class SkinGPT4(nn.Module):
         outputs = self.llama.generate(
             inputs_embeds=input_embeddings,
             max_new_tokens=max_new_tokens,
-            temperature=0.7,
+            temperature=0.3,
+            top_k=40,
             top_p=0.9,
             repetition_penalty=1.1,
             do_sample=True,
-            pad_token_id=self.tokenizer.eos_token_id
+            pad_token_id = self.tokenizer.eos_token_id,
+            eos_token_id = self.tokenizer.eos_token_id
         )