import os
os.system("pip install baseten")
os.system("pip install python-dotenv")
import gradio as gr
import baseten
from baseten.models import StableDiffusionPipeline
from dotenv import load_dotenv
load_dotenv()

BASETEN_API_KEY = os.getenv("BASETEN_API_KEY")
baseten.login(BASETEN_API_KEY)

model = StableDiffusionPipeline(
    "runwayml/stable-diffusion-v1-5",
    vae="stabilityai/sd-vae-ft-mse",
    scheduler="pndm"
)

def t2i(text):
    image, url = model(text, seed=2)
    return image

iface = gr.Interface(fn=t2i, inputs="text", outputs=[gr.Image(label="Generated Image")]).launch()
iface.launch(share=True)