fix added

SkinGPT.py

@@ -68,8 +68,10 @@ class Blip2QFormer(nn.Module):
         outputs = self.bert(
             attention_mask=attention_mask,
             inputs_embeds=combined_input,
+            output_attentions=True,
             return_dict=True
         )
+        plot_attention(outputs.attentions[-1][:, :, :32, 32:])
         return outputs.last_hidden_state[:, :self.num_query_tokens]
 
 
@@ -100,6 +102,15 @@ class SkinGPT4(nn.Module):
             self.q_former.bert_config.hidden_size,
             self.llama.config.hidden_size
         ).to(self.dtype)
+        self.tokenizer = LlamaTokenizer.from_pretrained(
+            "meta-llama/Llama-2-13b-chat-hf",
+            token=token,
+            padding_side="right"
+        )
+        self.tokenizer.add_special_tokens({
+            'additional_special_tokens': ['<IMAGE>']
+        })
+        self.llama.resize_token_embeddings(len(self.tokenizer))
         for module in [self.vit, self.ln_vision, self.q_former, self.llama_proj, self.llama]:
             for param in module.parameters():
                 param.requires_grad = False
@@ -170,30 +181,26 @@ class SkinGPT4(nn.Module):
             image_embeds = self.llama_proj(qformer_output.to(self.dtype))
         return image_embeds
 
-    def generate(self, images, user_input=None, max_new_tokens=300):
-        image_embeds = self.encode_image(images)
+    def generate_from_embeddings(self, image_embeds, user_input=None, max_new_tokens=300):
         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:
-        <IMAGE>
-        Could you describe the skin condition in this image?
-        ### Response:"""
-        self.tokenizer = LlamaTokenizer.from_pretrained(
-            "meta-llama/Llama-2-13b-chat-hf",
-            token=token,
-            padding_side="right"
-        )
-        num_added = self.tokenizer.add_special_tokens({
-            'additional_special_tokens': ['<IMAGE>']
-        })
-        if num_added == 0:
-            raise ValueError("Failed to add <IMAGE> token!")
-        self.llama.resize_token_embeddings(len(self.tokenizer))
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
+        
+        if user_input:
+            prompt = f"""### Instruction:
+            <IMAGE>
+            {user_input}
+            ### Response:"""
+        else:
+            prompt = """### Instruction:
+            <IMAGE>
+            Could you describe the skin condition in this image?
+            ### Response:"""
+
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(image_embeds.device)
         input_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
         visual_embeds = image_embeds.mean(dim=1)
         image_token_id = self.tokenizer.convert_tokens_to_ids("<IMAGE>")
@@ -218,7 +225,10 @@ class SkinGPT4(nn.Module):
         response = full_output.split("### Response:")[-1].strip()
         return response
 
-
+    def generate(self, images, user_input=None, max_new_tokens=300):
+        image_embeds = self.encode_image(images)
+        return self.generate_from_embeddings(image_embeds, user_input, max_new_tokens)
+    
 class SkinGPTClassifier:
     def __init__(self, device='cuda' if torch.cuda.is_available() else 'cpu'):
         self.device = torch.device(device)
@@ -229,6 +239,7 @@ class SkinGPTClassifier:
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
+        self.current_image_embeddings = None
 
     def _load_model(self):
         model_path = hf_hub_download(
@@ -239,11 +250,18 @@ class SkinGPTClassifier:
         model = model.to(self.device)
         return model
 
-    def predict(self, image):
-        image = image.convert('RGB')
-        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
+    def predict(self, image, user_input=None, reuse_embeddings=False):
+        if not reuse_embeddings or self.current_image_embeddings is None:
+            image = image.convert('RGB')
+            image_tensor = self.transform(image).unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                self.current_image_embeddings = self.model.encode_image(image_tensor)
+        
         with torch.no_grad():
-            diagnosis = self.model.generate(image_tensor)
+            diagnosis = self.model.generate_from_embeddings(
+                self.current_image_embeddings, 
+                user_input=user_input
+            )
         return {
             "diagnosis": diagnosis,
         }

app.py

@@ -81,11 +81,12 @@ uploaded_file = st.file_uploader(
 if uploaded_file is not None and uploaded_file != st.session_state.current_image:
     st.session_state.messages = []
     st.session_state.current_image = uploaded_file
+    classifier.current_image_embeddings = None
 
     image = Image.open(uploaded_file).convert("RGB")
     st.image(image, caption="Uploaded image", use_column_width=True)
     with st.spinner("Analyzing the image..."):
-        result = classifier.predict(image)
+        result = classifier.predict(image, reuse_embeddings=False)
 
     st.session_state.messages.append({"role": "assistant", "content": result["diagnosis"]})
 
@@ -101,18 +102,19 @@ if prompt := st.chat_input("Ask a follow-up question..."):
 
     with st.chat_message("assistant"):
         with st.spinner("Thinking..."):
+            image = Image.open(st.session_state.current_image).convert("RGB")
             if len(st.session_state.messages) > 1:
                 conversation_context = "\n".join(
                     f"{m['role']}: {m['content']}"
-                    for m in st.session_state.messages[:-1]  # Exclude current prompt
+                    for m in st.session_state.messages[:-1]
                 )
                 augmented_prompt = (
                     f"Conversation history:\n{conversation_context}\n\n"
                     f"Current question: {prompt}"
                 )
-                result = classifier.predict(image)
+                result = classifier.predict(image, user_input=augmented_prompt, reuse_embeddings=True)
             else:
-                result = classifier.predict(image)
+                result = classifier.predict(image, user_input=prompt, reuse_embeddings=False)
 
             st.markdown(result["diagnosis"])
             st.session_state.messages.append({"role": "assistant", "content": result["diagnosis"]})