Spaces:
Sleeping
Sleeping
File size: 3,501 Bytes
71c9483 1dcef5f 71c9483 1dcef5f 71c9483 1dcef5f 71c9483 e59f632 71c9483 e59f632 71c9483 e59f632 71c9483 |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import gradio as gr
import os
import torch
from transformers import AutoProcessor, MllamaForConditionalGeneration
from PIL import Image
import spaces
# Check if we're running in a Hugging Face Space and if SPACES_ZERO_GPU is enabled
IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
IS_SPACE = os.environ.get("SPACE_ID", None) is not None
# Determine the device (GPU if available, else CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"
LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
print(f"Using device: {device}")
print(f"Low memory mode: {LOW_MEMORY}")
# Get Hugging Face token from environment variables
HF_TOKEN = os.environ.get('HF_TOKEN')
# Load the model and processor
model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(
model_name,
use_auth_token=HF_TOKEN,
torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
device_map="auto" if device == "cuda" else None, # Use device mapping if CUDA is available
)
# Move the model to the appropriate device (GPU if available)
model.to(device)
processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN)
@spaces.GPU # Use the free GPU provided by Hugging Face Spaces
def predict(image, text):
# Prepare the input messages
messages = [
{"role": "user", "content": [
{"type": "image"}, # Specify that an image is provided
{"type": "text", "text": text} # Add the user-provided text input
]}
]
# Create the input text using the processor's chat template
input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
# Process the inputs and move to the appropriate device
inputs = processor(image, input_text, return_tensors="pt").to(device)
# Generate a response from the model
outputs = model.generate(**inputs, max_new_tokens=100)
# Decode the output to return the final response
response = processor.decode(outputs[0], skip_special_tokens=True)
return response
def predict_text(text):
# Prepare the input messages
messages = [{"role": "user", "content": [{"type": "text", "text": txt}]}]
# Create the input text using the processor's chat template
input_text = processor.apply_chat_template(messages, add_generation_prompt=True)
# Process the inputs and move to the appropriate device
# inputs = processor(image, input_text, return_tensors="pt").to(device)
inputs = processor(text=text, return_tensors="pt").to("cuda")
# Generate a response from the model
outputs = model.generate(**inputs, max_new_tokens=250)
# Decode the output to return the final response
response = processor.decode(outputs[0], skip_special_tokens=True)
return response
# Define the Gradio interface
interface = gr.Interface(
fn=predict_text,
inputs=[
# gr.Image(type="pil", label="Image Input"), # Image input with label
gr.Textbox(label="Text Input") # Textbox input with label
],
outputs=gr.Textbox(label="Generated Response"), # Output with a more descriptive label
title="Llama 3.2 11B Vision Instruct Demo", # Title of the interface
description="This demo uses Meta's Llama 3.2 11B Vision model to generate responses based on an image and text input.", # Short description
theme="compact" # Using a compact theme for a cleaner look
)
# Launch the interface
interface.launch(debug=True) |