Issue fix

app.py

@@ -340,15 +340,16 @@ class SkinGPT4(nn.Module):
         return prompt
 
     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)
         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,
@@ -357,7 +358,7 @@ class SkinGPT4(nn.Module):
             top_p=0.9,
             do_sample=True
         )
-
+        print("Generated diagnosis")
         return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 class SkinGPTClassifier:
@@ -430,12 +431,14 @@ if uploaded_file:
     st.image(uploaded_file, caption="Uploaded image", use_column_width=True)
     image = Image.open(uploaded_file).convert("RGB")
     if not st.session_state.conversation:
-        # First message - diagnosis
         with st.spinner("Analyzing image..."):
-            diagnosis = classifier.predict(image)
-        st.session_state.conversation.append(("assistant", diagnosis))
-        with st.chat_message("assistant"):
-            st.markdown(diagnosis)
+            result = classifier.predict(image)
+            if "error" in result:
+                st.error(result["error"])
+            else:
+                st.session_state.conversation.append(("assistant", result))
+                with st.chat_message("assistant"):
+                    st.markdown(result)
     else:
         # Follow-up questions
         if user_query := st.chat_input("Ask a follow-up question..."):