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

Dir = os.path.dirname(os.path.abspath(__file__))
file = f"{Dir}/1.png"
if not os.path.isfile(file):
    open(file, "wb").write(b"")

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 greet1(inp):
    if inp:
        u = agent.run("generate an image of `text` ", answer=inp)
        u.save(file)
        return Image.open(file)

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