fix added

SkinGPT.py

@@ -22,6 +22,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 class Blip2QFormer(nn.Module):
     def __init__(self, num_query_tokens=32, vision_width=1408):
         super().__init__()
+        self.num_query_tokens = num_query_tokens
         # Load pre-trained Q-Former config
         self.bert_config = BertConfig(
             vocab_size=30522,
@@ -74,26 +75,24 @@ class Blip2QFormer(nn.Module):
         visual_embeds = self.vision_proj(visual_features.float())
         print(f"Projected embeds stats - min: {visual_embeds.min().item():.4f}, max: {visual_embeds.max().item():.4f}")
         # visual_embeds = self.vision_proj(visual_features.float())
-        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)
+        combined_input = torch.cat([query_tokens, visual_embeds], dim=1)
+        attention_mask = torch.ones(
+            combined_input.size()[:-1],
+            dtype=torch.long,
+            device=combined_input.device
+        )
 
         # Forward through BERT
         outputs = self.bert(
-            input_ids=None,  # No text input
-            attention_mask=None,
-            inputs_embeds=query_tokens,
-            encoder_hidden_states=visual_embeds,
-            encoder_attention_mask=visual_attention_mask,
+            attention_mask=attention_mask,
+            inputs_embeds=combined_input,
             return_dict=True
         )
 
-        return outputs.last_hidden_state
+        return outputs.last_hidden_state[:, :self.num_query_tokens]