File size: 2,186 Bytes
ae9601e
a225caf
2064a53
 
a225caf
 
08b3c07
a225caf
 
 
ae9601e
 
 
 
 
 
 
 
a225caf
ae9601e
257403d
 
 
 
 
 
a225caf
 
 
 
 
 
 
2064a53
 
a225caf
 
 
 
 
 
 
 
2064a53
 
 
 
 
a225caf
 
 
 
 
 
 
 
 
 
 
 
 
 
257403d
 
a225caf
 
 
 
2064a53
a225caf
2064a53
a225caf
 
 
 
 
2064a53
 
a225caf
 
ae9601e
a225caf
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
from gradio_client import Client, handle_file
import gradio as gr
import os


MODELS = {
    "SmolVLM-Instruct": "akhaliq/SmolVLM-Instruct"
}

def create_chat_fn(client):
    def chat(message, history):
        # Extract text and files from the message
        text = message.get("text", "")
        files = message.get("files", [])
        
        # Handle file uploads if present
        processed_files = [handle_file(f) for f in files]
        
        response = client.predict(
            message={"text": text, "files": processed_files},
            system_prompt="You are a helpful AI assistant.",
            temperature=0.7,
            max_new_tokens=1024,
            top_k=40,
            repetition_penalty=1.1,
            top_p=0.95,
            api_name="/chat"
        )
        return response
    return chat

def set_client_for_session(model_name, request: gr.Request):
    headers = {}
    if request and hasattr(request, 'headers'):
        x_ip_token = request.headers.get('x-ip-token')
        if x_ip_token:
            headers["X-IP-Token"] = x_ip_token
    
    return Client(MODELS[model_name], headers=headers)

def safe_chat_fn(message, history, client):
    if client is None:
        return "Error: Client not initialized. Please refresh the page."
    try:
        return create_chat_fn(client)(message, history)
    except Exception as e:
        print(f"Error during chat: {str(e)}")
        return f"Error during chat: {str(e)}"

with gr.Blocks() as demo:
    
    client = gr.State()
    
    model_dropdown = gr.Dropdown(
        choices=list(MODELS.keys()),
        value="SmolVLM-Instruct",
        label="Select Model",
        interactive=True
    )
    
    chat_interface = gr.ChatInterface(
        fn=safe_chat_fn,
        additional_inputs=[client],
        multimodal=True
    )
    
    # Update client when model changes
    model_dropdown.change(
        fn=set_client_for_session,
        inputs=[model_dropdown],
        outputs=[client]
    )
    
    # Initialize client on page load
    demo.load(
        fn=set_client_for_session,
        inputs=[gr.State("SmolVLM-Instruct")],
        outputs=[client]
    )

demo.launch()