import gradio as gr
import spaces
import torch
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline,AutoPipelineForText2Image, DEISMultistepScheduler

pipe = AutoPipelineForText2Image.from_pretrained('Lykon/AAM_XL_AnimeMix', torch_dtype=torch.float16, variant="fp16")
pipe.scheduler = DEISMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

@spaces.GPU
def predict(prompt):
  return pipe(prompt, num_inference_steps=25).images[0]

demo = gr.Interface(
  fn=predict,
  inputs='text',
  outputs= 'image',
)
demo.launch()