import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

model_id = "Fazzie/PokemonGAI"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
print("load model successfully!")
pipe = pipe.to(device)

def imagen(inputs):

    prompt = inputs
    negative_prompt = 'pixelated, distorted'

    print("start inference")
    image = pipe(prompt=inputs, negative_prompt=negative_prompt).images[0]
    print("finish inference!")
    return image
    
demo = gr.Interface(fn=imagen, inputs=[gr.Textbox(label = 'Prompt')], outputs="image")

if __name__ == "__main__":
    demo.launch()

image.save("1.png")

# gr.Interface.load("models/Fazzie/PokemonGAI").launch()