import gradio as gr

# Load the model
model_interface = gr.load("models/google/gemma-2-2b-it")

# Create a Gradio interface with custom title and logo
def main():
    with gr.Blocks() as demo:
        # Add the title
        gr.Markdown("# Gemma 2b Demo")
        
        # Add the logo
        gr.Image("orYx logo.png", elem_id="logo", show_label=False, interactive=False)
        
        # Embed the model interface inside a Row
        with gr.Row():
            model_interface.render()  # Use render to include the prebuilt interface

    demo.launch()

if __name__ == "__main__":
    main()