app_huggingface.py

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