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

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])

if __name__ == "__main__":
    demo.launch()