Spaces:
Runtime error
Runtime error
File size: 2,508 Bytes
5d75401 5c78d10 5d75401 5c78d10 5d75401 5c78d10 5d75401 5c78d10 5d75401 5c78d10 5d75401 5c78d10 5d75401 5c78d10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPProcessor, CLIPModel
from PIL import Image
import requests
from io import BytesIO
# 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")
# 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)
# 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 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)
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
demo.launch()
|