Issue fix

app.py

@@ -54,22 +54,18 @@ class Blip2QFormer(nn.Module):
         )
 
         self.bert = BertModel(self.bert_config, add_pooling_layer=False).to(torch.float16)
-
-        # Replace position embeddings with a dummy implementation
-        self.bert.embeddings.position_embeddings = nn.Identity()  # Completely bypass position embeddings
-
-        # Disable word embeddings
-        self.bert.embeddings.word_embeddings = None
-
-        # Initialize query tokens
         self.query_tokens = nn.Parameter(
-            torch.zeros(1, num_query_tokens, self.bert_config.hidden_size, dtype=torch.float16)
+            torch.zeros(1, num_query_tokens, self.bert_config.hidden_size)
         )
-        self.vision_proj = nn.Sequential(
-            nn.Linear(vision_width, self.bert_config.hidden_size),
-            nn.LayerNorm(self.bert_config.hidden_size)
-        ).to(torch.float16)
+        self.vision_proj = nn.Linear(vision_width, self.bert_config.hidden_size)
+
+        # Initialize weights
+        self._init_weights()
 
+    def _init_weights(self):
+        nn.init.normal_(self.query_tokens, std=0.02)
+        nn.init.xavier_uniform_(self.vision_proj.weight)
+        nn.init.constant_(self.vision_proj.bias, 0)
 
     def load_from_pretrained(self, url_or_filename):
         if url_or_filename.startswith('http'):
@@ -77,38 +73,31 @@ class Blip2QFormer(nn.Module):
             checkpoint = torch.load(BytesIO(response.content), map_location='cpu')
         else:
             checkpoint = torch.load(url_or_filename, map_location='cpu')
-
-        # Load Q-Former weights only
         state_dict = checkpoint['model'] if 'model' in checkpoint else checkpoint
         msg = self.load_state_dict(state_dict, strict=False)
-        # print(f"Loaded Q-Former weights with message: {msg}")
-
-    def forward(self, query_embeds=None, encoder_hidden_states=None, encoder_attention_mask=None):
-        if query_embeds is None:
-            query_embeds = self.query_tokens.expand(encoder_hidden_states.shape[0], -1, -1)
 
+    def forward(self, visual_features):
         # Project visual features
-        visual_embeds = self.vision_proj(encoder_hidden_states)
-
-        # Create proper attention mask
-        if encoder_attention_mask is None:
-            encoder_attention_mask = torch.ones(
-                visual_embeds.size()[:-1],
-                dtype=torch.long,
-                device=visual_embeds.device
-            )
-        batch_size = query_embeds.size(0)
-        extended_attention_mask = encoder_attention_mask.unsqueeze(1).expand(-1, query_embeds.size(1), -1)
+        visual_embeds = self.vision_proj(visual_features)
+        visual_attention_mask = torch.ones(
+            visual_embeds.size()[:-1],
+            dtype=torch.long,
+            device=visual_embeds.device
+        )
+
+        # Expand query tokens
+        query_tokens = self.query_tokens.expand(visual_embeds.shape[0], -1, -1)
 
-        encoder_outputs = self.bert.encoder(
-            hidden_states=query_embeds,
+        # Forward through BERT
+        outputs = self.bert(
+            None,  # No text input
             attention_mask=None,
             encoder_hidden_states=visual_embeds,
-            encoder_attention_mask=encoder_attention_mask,
+            encoder_attention_mask=visual_attention_mask,
             return_dict=True
         )
-        return encoder_outputs.last_hidden_state
 
+        return outputs.last_hidden_state
 
 class LayerNorm(nn.LayerNorm):
     """Subclass torch's LayerNorm to handle fp16."""
@@ -137,19 +126,13 @@ class SkinGPT4(nn.Module):
                  q_former_model="https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth"):
         super().__init__()
         # Image encoder parameters from paper
+        self.device = device
         self.dtype = torch.float16
         self.H, self.W, self.C = 224, 224, 3
         self.P = 14  # Patch size
         self.D = 1408  # ViT embedding dimension
         self.num_query_tokens = 32
 
-        # self.tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf",
-        #                                                 token=token, padding_side="right")
-        #
-        # print("Loaded tokenizer")
-        # self.tokenizer.add_special_tokens({'additional_special_tokens': ['<ImageHere>']})
-
-        # Initialize components
         self.vit = self._init_vit(vit_checkpoint_path)
         print("Loaded ViT")
         self.ln_vision = nn.LayerNorm(self.D).to(self.dtype)
@@ -161,7 +144,10 @@ class SkinGPT4(nn.Module):
         self.q_former.load_from_pretrained(q_former_model)
         for param in self.q_former.parameters():
             param.requires_grad = False
-        self.q_former.eval()
+        for module in [self.vit, self.ln_vision, self.q_former]:
+            for param in module.parameters():
+                param.requires_grad = False
+            module.eval()
         print("Loaded QFormer")
 
         self.tokenizer = LlamaTokenizer.from_pretrained(
@@ -169,24 +155,18 @@ class SkinGPT4(nn.Module):
             token=token,
             padding_side="right"
         )
-        self.tokenizer.add_special_tokens({'additional_special_tokens': ['<Img>', '</Img>', '<Image>']})
+        self.tokenizer.add_special_tokens({'additional_special_tokens': ['<Img>', '</Img>', '<ImageHere>']})
 
         self.llama = self._init_llama()
-        # self.llama.resize_token_embeddings(len(self.tokenizer))
         self.llama.resize_token_embeddings(len(self.tokenizer))
 
         self.llama_proj = nn.Linear(
             self.q_former.bert_config.hidden_size,
             self.llama.config.hidden_size
         ).to(self.dtype)
-        self._init_alignment_projection()
-        print("Loaded Llama")
-        # Initialize learnable query tokens
 
-        self.query_tokens = nn.Parameter(
-            torch.zeros(1, self.num_query_tokens, self.q_former.bert_config.hidden_size)
-        )
-        nn.init.normal_(self.query_tokens, std=0.02)
+        for param in self.llama_proj.parameters():
+            param.requires_grad = False
 
     def _init_vit(self, vit_checkpoint_path):
         """Initialize EVA-ViT-G with paper specifications"""
@@ -297,28 +277,13 @@ class SkinGPT4(nn.Module):
         return x  # (B, N+1, D)
 
     def forward(self, images):
-        images = images.to(self.dtype)
-        # Convert images to patches
-        x = self._create_patches(images)  # (B, N+1, D)
-
-        # ViT processing
-        x = x.to(self.dtype)
-        self.vit = self.vit.to(self.dtype)
-        vit_output = self.forward_encoder(x)  # (B, N+1, D)
-
-        # Q-Former processing
-        query_tokens = self.query_tokens.expand(x.size(0), -1, -1).to(torch.float16)
-        qformer_output = self.q_former(
-            query_embeds=query_tokens,
-            encoder_hidden_states=vit_output.to(torch.float16),
-            encoder_attention_mask=torch.ones_like(vit_output[:, :, 0])
-        ).to(self.dtype)
-
-        # Alignment projection
-        aligned_features = self.llama_proj(qformer_output.to(self.dtype))
-
+        x = self._create_patches(images)
+        vit_output = self.forward_encoder(x)
+        qformer_output = self.q_former(vit_output)
+        aligned_features = self.llama_proj(qformer_output)
         return aligned_features
 
+
     def add_to_history(self, role, content):
         self.conversation_history.append({"role": role, "content": content})
 
@@ -347,85 +312,42 @@ class SkinGPT4(nn.Module):
 
     def generate(self, images, user_input=None, max_length=300):
         print("Analysing the image to generate the diagnosis")
-        # Get aligned features
         aligned_features = self.forward(images)
         print("Generated the aligned features with ViT and Qformer")
-        # prompt = self.build_prompt(aligned_features, user_input)
-        # prompt = (
-        #     "### Instruction: <Img><ImageHere></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:"
+        prompt = (
+            "[INST] <<SYS>>\n"
+            "You are a dermatology AI assistant. Analyze this skin image carefully and provide:\n"
+            "1. A description of visible features\n"
+            "2. Potential diagnoses\n"
+            "3. Recommendations for next steps\n"
+            "<</SYS>>\n\n"
+            "<Img><ImageHere></Img> Could you describe the skin disease in this image for me? [/INST]"
+        )
         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)
-
-        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]}")
-
+        image_token_id = self.tokenizer.convert_tokens_to_ids("<ImageHere>")
+        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")
         # 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)
+        input_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
+        projected_features = self.llama_proj(aligned_features.mean(dim=1, keepdim=True))
+        input_embeddings[image_token_pos] = projected_features
 
         outputs = self.llama.generate(
-            inputs_embeds=embeddings,
-            max_length=max_length,
+            inputs_embeds=input_embeddings,
+            max_new_tokens=max_length,
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
-            pad_token_id=self.tokenizer.eos_token_id
+            pad_token_id=self.tokenizer.eos_token_id,
+            attention_mask=inputs.attention_mask,
+            num_return_sequences=1
         )
+
         print(f"Output from llama : {outputs}")
         full_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response_start = full_output.find("### Response:") + len("### Response:")
-        return full_output[response_start:].strip()
-
-
-        # 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': ['<ImageHere>']})
-        # self.llama.resize_token_embeddings(len(self.tokenizer))
-        # tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf",
-        #                                            token=token, padding_side="right")
-        # tokenizer.add_special_tokens({'additional_special_tokens': ['<ImageHere>']})
-        # self.llama.resize_token_embeddings(len(tokenizer))
-        # inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
-        # image_embeddings = self.llama.model.embed_tokens(inputs.input_ids)
-        # image_token_index = torch.where(inputs.input_ids == self.tokenizer.convert_tokens_to_ids("<ImageHere>"))
-        # image_embeddings[image_token_index] = aligned_features.mean(dim=1)  # Pool query tokens
-        # print("Generating the diagnosis with llama")
-        # # Generate response
-        # outputs = self.llama.generate(
-        #     inputs_embeds=image_embeddings,
-        #     max_length=max_length,
-        #     temperature=0.7,
-        #     top_p=0.9,
-        #     do_sample=True
-        # )
-        # print("Generated diagnosis")
-        # return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return full_output.split("[/INST]")[-1].strip()
+
 
 class SkinGPTClassifier:
     def __init__(self, device='cuda' if torch.cuda.is_available() else 'cpu'):
@@ -433,8 +355,7 @@ class SkinGPTClassifier:
         self.conversation_history = []
 
         with st.spinner("Loading AI models (this may take several minutes)..."):
-            self.meta_model = self.load_models()
-        self.resnet_feature_extractor = None
+            self.model = self._load_model()
 
         # Image transformations
         self.transform = transforms.Compose([
@@ -443,26 +364,27 @@ class SkinGPTClassifier:
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
 
-    def load_models(self):
+    def _load_model(self):
         model_path = hf_hub_download(
             repo_id="KeerthiVM/SkinCancerDiagnosis",
             filename="dermnet_finetuned_version1.pth",
         )
-        meta_model = SkinGPT4(vit_checkpoint_path=model_path)
-        return meta_model
+        model = SkinGPT4(vit_checkpoint_path=model_path).eval()
+        model = model.to(self.device)
+        model.eval()
+        return model
 
     def predict(self, image):
         image = image.convert('RGB')
         image_tensor = self.transform(image).unsqueeze(0).to(self.device)
-        diagnosis = self.meta_model.generate(
-            image_tensor
-        )
+        with torch.no_grad():
+            diagnosis = self.model.generate(image_tensor)
 
         return {
-            "top_predictions": diagnosis,
+            "diagnosis": diagnosis,
+            "visual_features": None  # Can return features if needed
         }
 
-# @st.cache_resource
 def get_classifier():
     return SkinGPTClassifier()