Issue fix

app.py

@@ -24,6 +24,14 @@ from transformers import BitsAndBytesConfig
 from accelerate import init_empty_weights
 import warnings
 from transformers import logging
+
+import torch
+from torch.cuda.amp import autocast
+
+# Set default dtypes
+torch.set_default_dtype(torch.float32)  # Main computations in float32
+MODEL_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
+
 warnings.filterwarnings("ignore", category=FutureWarning, module="timm")
 warnings.filterwarnings("ignore", category=UserWarning)
 logging.set_verbosity_error()
@@ -83,7 +91,9 @@ class Blip2QFormer(nn.Module):
 
     def forward(self, visual_features):
         # Project visual features
-        visual_embeds = self.vision_proj(visual_features.float())
+        with autocast(enabled=False):
+            visual_embeds = self.vision_proj(visual_features.float())
+        # visual_embeds = self.vision_proj(visual_features.float())
         visual_attention_mask = torch.ones(
             visual_embeds.size()[:-1],
             dtype=torch.long,
@@ -133,20 +143,21 @@ class SkinGPT4(nn.Module):
         super().__init__()
         # Image encoder parameters from paper
         self.device = device
-        self.dtype = torch.float16
+        # self.dtype = torch.float16
+        self.dtype = MODEL_DTYPE
         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.vit = self._init_vit(vit_checkpoint_path)
+        self.vit = self._init_vit(vit_checkpoint_path).to(self.dtype)
         print("Loaded ViT")
         self.ln_vision = nn.LayerNorm(self.D).to(self.dtype)
 
         self.q_former = Blip2QFormer(
             num_query_tokens=self.num_query_tokens,
             vision_width=self.D
-        ).to(self.dtype)
+        )
         self.q_former.load_from_pretrained(q_former_model)
         for param in self.q_former.parameters():
             param.requires_grad = False
@@ -368,8 +379,8 @@ class SkinGPTClassifier:
 
         with st.spinner("Loading AI models (this may take several minutes)..."):
             self.model = self._load_model()
-        print(f"Q-Former output shape: {self.model.q_former(torch.randn(1, 197, 1408)).shape}")
-        print(f"Projection layer: {self.model.llama_proj}")
+        # print(f"Q-Former output shape: {self.model.q_former(torch.randn(1, 197, 1408)).shape}")
+        # print(f"Projection layer: {self.model.llama_proj}")
 
         # Image transformations
         self.transform = transforms.Compose([