Issue fix

app.py

@@ -24,6 +24,7 @@ from transformers import BitsAndBytesConfig
 from accelerate import init_empty_weights
 import warnings
 from transformers import logging
+warnings.filterwarnings("ignore", category=FutureWarning, module="timm")
 warnings.filterwarnings("ignore", category=UserWarning)
 logging.set_verbosity_error()
 token = os.getenv("HF_TOKEN")
@@ -57,7 +58,7 @@ class Blip2QFormer(nn.Module):
             classifier_dropout=None,
         )
 
-        self.bert = BertModel(self.bert_config, add_pooling_layer=False).to(torch.float16)
+        self.bert = BertModel(self.bert_config, add_pooling_layer=False)
         self.query_tokens = nn.Parameter(
             torch.zeros(1, num_query_tokens, self.bert_config.hidden_size)
         )
@@ -82,7 +83,7 @@ class Blip2QFormer(nn.Module):
 
     def forward(self, visual_features):
         # Project visual features
-        visual_embeds = self.vision_proj(visual_features)
+        visual_embeds = self.vision_proj(visual_features.float())
         visual_attention_mask = torch.ones(
             visual_embeds.size()[:-1],
             dtype=torch.long,
@@ -285,10 +286,12 @@ class SkinGPT4(nn.Module):
         return x  # (B, N+1, D)
 
     def forward(self, images):
+        images = images.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)
+        with torch.cuda.amp.autocast(enabled=False):
+            qformer_output = self.q_former(vit_output.float())
+        aligned_features = self.llama_proj(qformer_output.to(self.dtype))
         return aligned_features