Removed secrets from history

app.py

@@ -20,7 +20,8 @@ import requests
 from io import BytesIO
 import os
 from huggingface_hub import hf_hub_download
-
+from transformers import BitsAndBytesConfig
+from accelerate import init_empty_weights
 
 token = os.getenv("HF_TOKEN")
 if not token:
@@ -164,7 +165,6 @@ class SkinGPT4(nn.Module):
         self.q_former.eval()
         print("Loaded QFormer")
         self.llama = self._init_llama()
-        self.llama = self.llama.to(device)
         self.llama.resize_token_embeddings(len(self.tokenizer))
 
         self.llama_proj = nn.Linear(
@@ -214,30 +214,16 @@ class SkinGPT4(nn.Module):
     def _init_llama(self):
         """Initialize frozen LLaMA-2-13b-chat with proper error handling"""
         try:
-            from transformers import BitsAndBytesConfig
-            from accelerate import init_empty_weights
-
-            # Configure 4-bit quantization to reduce memory usage
-            # quantization_config = BitsAndBytesConfig(
-            #     load_in_4bit=True,
-            #     bnb_4bit_compute_dtype=torch.float16,
-            #     bnb_4bit_use_double_quant=True,
-            #     bnb_4bit_quant_type="nf4"
-            # )
-            quant_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16,
-                bnb_4bit_quant_type="nf4",
-            )
-
+            device_map = {
+                "": 0 if torch.cuda.is_available() else "cpu"
+            }
             # First try loading with device_map="auto"
             try:
                 model = LlamaForCausalLM.from_pretrained(
                     "meta-llama/Llama-2-13b-chat-hf",
-                    # quantization_config=quant_config,
                     token=token,
                     torch_dtype=torch.float16,
-                    device_map="auto",
+                    device_map=device_map,
                     low_cpu_mem_usage=True
                 )
             except ImportError:
@@ -355,22 +341,10 @@ class SkinGPT4(nn.Module):
 
     def generate(self, images, user_input=None, max_length=300):
         # Get aligned features
-        images = images.to(self.dtype)
-
         aligned_features = self.forward(images)
 
         prompt = self.build_prompt(aligned_features, user_input)
-
-        self.llama = self.llama.to(self.dtype)
-
-        # Tokenize prompt
-
-        # self.tokenizer.add_special_tokens({'additional_special_tokens': ['<ImageHere>']})
-        # self.llama.resize_token_embeddings(len(self.tokenizer))
-
         inputs = self.tokenizer(prompt, return_tensors="pt").to(images.device)
-
-        # Replace <ImageHere> with aligned features
         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
@@ -386,27 +360,13 @@ class SkinGPT4(nn.Module):
 
         return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-
-# def load_model(model_path):
-#     model_path = hf_hub_download(
-#         repo_id="KeerthiVM/SkinCancerDiagnosis",
-#         filename="dermnet_finetuned_version1.pth",
-#     )
-#     # model = SkinGPT4(vit_checkpoint_path="dermnet_finetuned_version1.pth")
-#     model = SkinGPT4(vit_checkpoint_path=model_path)
-#     model.to(device)
-#     model.eval()
-#     return model
-
-
-
 class SkinGPTClassifier:
     def __init__(self, device='cuda' if torch.cuda.is_available() else 'cpu'):
         self.device = torch.device(device)
         self.conversation_history = []
-        # Initialize models (they'll be loaded when needed)
-        self.base_models = None
-        self.meta_model = None
+
+        with st.spinner("Loading AI models (this may take several minutes)..."):
+            self.meta_model = self.load_models()
         self.resnet_feature_extractor = None
 
         # Image transformations
@@ -421,22 +381,11 @@ class SkinGPTClassifier:
             repo_id="KeerthiVM/SkinCancerDiagnosis",
             filename="dermnet_finetuned_version1.pth",
         )
-        # self.meta_model = SkinGPT4(vit_checkpoint_path="dermnet_finetuned_version1.pth")
-        self.meta_model = SkinGPT4(vit_checkpoint_path=model_path)
-        self.meta_model.to_empty(device=device)
-
-    def predict(self, image, top_k=3):
-        """Make prediction for a single image"""
-        if self.meta_model is None:
-            self.load_models()
-
-        # Load and preprocess image
-        try:
-            # image = Image.open(image_path).convert('RGB')
-            image = image.convert('RGB')
-        except:
-            raise ValueError("Could not load image from path")
+        meta_model = SkinGPT4(vit_checkpoint_path=model_path)
+        return meta_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
@@ -446,18 +395,16 @@ class SkinGPTClassifier:
             "top_predictions": diagnosis,
         }
 
-classifier = SkinGPTClassifier()
+@st.cache_resource
+def get_classifier():
+    return SkinGPTClassifier()
 
+classifier = get_classifier()
 
 # === Session Init ===
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# === Image Processing Function ===
-def run_inference(image):
-    result = classifier.predict(image, top_k=1)
-
-    return result
 
 # === PDF Export ===
 def export_chat_to_pdf(messages):
@@ -484,7 +431,8 @@ if uploaded_file:
     image = Image.open(uploaded_file).convert("RGB")
     if not st.session_state.conversation:
         # First message - diagnosis
-        diagnosis =  classifier.predict(image, top_k=1)
+        with st.spinner("Analyzing image..."):
+            diagnosis = classifier.predict(image)
         st.session_state.conversation.append(("assistant", diagnosis))
         with st.chat_message("assistant"):
             st.markdown(diagnosis)