import gradio as gr
import os
import torch
from transformers import pipeline
from diffusers import StableDiffusionPipeline
from gradio.components import Image
from gradio.components import Textbox



model_id = "CompVis/stable-diffusion-v1-4"
my_model = "waterplayfire/MyModel"
device = "cuda"

cuda_support = torch.cuda.is_available()
title = "Txt to Image"
pipeline = None

if cuda_support:
    title = "Txt to Image by GPU"
    pipeline = StableDiffusionPipeline.from_pretrained(my_model, torch_dtype=torch.float16, revision="fp16")
    pipeline = pipeline.to(device)
    torch.backends.cudnn.enabled = True
else:
    title = "Txt to Image by CPU"
    pipeline = StableDiffusionPipeline.from_pretrained(my_model)

#title = os.environ['model_fetch']

def predict(prompt):
  predictions = pipeline(prompt, guidance_scale=7.5).images[0]
  return predictions#{p["label"]: p["score"] for p in predictions}

gr.Interface(
    predict,
    inputs=Textbox(),
    outputs=Image(),
    title=title,
    layout="vertical",
).launch()