Update ig.py

ig.py

@@ -1,10 +1,21 @@
 from transformers import AutoProcessor, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
 import gradio as gr
 import torch
 
-# Load the processor and model
+# Load the processor
 processor = AutoProcessor.from_pretrained("microsoft/git-base")
-model = AutoModelForCausalLM.from_pretrained("./")
+
+# Load the base model (the pre-trained model you're adapting with LoRA)
+base_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base")
+
+# Load the adapter configuration
+adapter_config_path = "./"  # Path to your adapter_config.json
+adapter_model_path = "./"  # Path to your adapter_model.safetensors
+
+# Load the LoRA adapter using Peft
+peft_config = PeftConfig.from_pretrained(adapter_config_path)
+model = PeftModel.from_pretrained(base_model, adapter_model_path, config=peft_config)
 
 def predict(image):
     try:
@@ -28,7 +39,7 @@ def predict(image):
         print("Error during prediction:", str(e))
         return "Error: " + str(e)
 
-# https://www.gradio.app/guides
+# Gradio Interface
 with gr.Blocks() as demo:
     image = gr.Image(type="pil")
     predict_btn = gr.Button("Predict", variant="primary")