Update app.py

app.py

@@ -1,40 +1,55 @@
 import gradio as gr
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPProcessor, CLIPModel
 from PIL import Image
 import requests
 from io import BytesIO
 
-# Load YOLOv5 pre-trained model from Hugging Face
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # You can choose other versions like yolov5m or yolov5l
+# Load CLIP model for image classification
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-# Function for object detection
-def detect_objects(input_image):
-    # If the input is a URL, download the image
+# Load Mistral-7B-Instruct-v0.3 model for chat
+mistral_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+mistral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+
+# Function for image classification with CLIP (anime recognition)
+def classify_image(input_image):
     if isinstance(input_image, str):
         response = requests.get(input_image)
         img = Image.open(BytesIO(response.content))
     else:
         img = Image.fromarray(input_image)
     
-    # Run YOLOv5 object detection
-    results = model(img)
-    
-    # Render results on image
-    results.render()  # Render boxes on the image
+    # Prepare the image and text (anime-related labels)
+    inputs = clip_processor(text=["anime", "cartoon", "realistic", "painting"], images=img, return_tensors="pt", padding=True)
+    outputs = clip_model(**inputs)
+    logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+    probs = logits_per_image.softmax(dim=1)  # we can apply softmax to get the label probabilities
     
-    # Return image with detections
-    output_image = results.imgs[0]
-    return Image.fromarray(output_image)
+    # Return the predicted class label
+    labels = ["anime", "cartoon", "realistic", "painting"]
+    predicted_label = labels[probs.argmax()]
+    return predicted_label
+
+# Function for chat with Mistral 7B Instruct
+def chat_with_mistral(input_text):
+    inputs = mistral_tokenizer(input_text, return_tensors="pt")
+    outputs = mistral_model.generate(inputs["input_ids"], max_length=150)
+    response = mistral_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# Create Gradio interface for both Image Classification and Chat
+with gr.Blocks() as demo:
+    with gr.Tab("Chat with Mistral"):
+        chat_input = gr.Textbox(label="Ask Mistral 7B", placeholder="Type your question here...")
+        chat_output = gr.Textbox(label="Mistral's Reply", interactive=False)
+        chat_input.submit(chat_with_mistral, inputs=chat_input, outputs=chat_output)
 
-# Create Gradio interface
-interface = gr.Interface(
-    fn=detect_objects, 
-    inputs=gr.inputs.Image(type="numpy", label="Upload an image"), 
-    outputs=gr.outputs.Image(type="pil", label="Detected Image"), 
-    title="YOLOv5 Object Detection",
-    description="Upload an image and detect objects using YOLOv5 model. The model can identify objects like people, cars, animals, and more.",
-    theme="huggingface"
-)
+    with gr.Tab("Classify Anime Image"):
+        img_input = gr.Image(type="numpy", label="Upload Image for Anime Classification")
+        img_output = gr.Textbox(label="Predicted Label", interactive=False)
+        img_input.change(classify_image, inputs=img_input, outputs=img_output)
 
 # Launch the interface
-interface.launch()
+demo.launch()