Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from PIL import Image
-from vit_model_test import CustomModel  # Ensure this is the correct import for your model
+from vit_model_test import CustomModel  
 from vit_Training import Custom_VIT_Model
 
 # Initialize the model
@@ -32,7 +32,7 @@ custom_css = """
 </style>
 """
 
-# Variable to store the last prediction result
+# Store the last prediction result
 last_prediction = None
 
 def predict(image: Image.Image):
@@ -80,7 +80,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         submit_btn = gr.Button("Submit", variant="primary", elem_id="submit_btn")
         feedback_btn = gr.Button("The model was wrong", variant="secondary", elem_id="feedback_btn")
         
-    gr.Markdown("<br>")  # Adding space below buttons
+    gr.Markdown("<br>")  
     clear_btn = gr.Button("Clear", elem_id="clear_btn")  
 
     submit_btn.click(predict, inputs=image_input, outputs=[prediction_output, confidence_output])
@@ -88,7 +88,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 
     # Clear button 
     def clear_all():
-        return None, "", "", ""  # Clear all outputs
+        return None, "", "", ""  
 
     clear_btn.click(clear_all, outputs=[image_input, prediction_output, confidence_output, feedback_output])