import gradio as gr
import spaces
from transformers import pipeline

# Load a small model that runs fast
generator = pipeline("text-generation", model="gpt2")

@spaces.GPU
def generate(text):
    result = generator(text, max_length=50, num_return_sequences=1)
    return result[0]['generated_text']

# Create the UI
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 My Free H200 Space")
    input_text = gr.Textbox(label="Enter a prompt")
    output_text = gr.Textbox(label="Generated Text")
    btn = gr.Button("Generate")
    
    btn.click(fn=generate, inputs=input_text, outputs=output_text)

demo.launch()