import gradio as gr
import torch
import os
from PIL import Image

from transformers import AutoModelForCausalLM, AutoTokenizer, LocalAgent

checkpoint = "cerebras/Cerebras-GPT-1.3B"
agent = LocalAgent.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True)

def greet(inp):
    if inp:
        return agent.run("generate an image of `text` ", answer=inp)

iface = gr.Interface(fn=greet, inputs="text", outputs="image")
iface.launch()