import gradio as gr
from diffusers import LDMTextToImagePipeline
import torch
import numpy as np
import PIL
import cv2

import PIL.Image
import random

print('\nDEBUG:   Version: 3')

#pipeline = LDMTextToImagePipeline.from_pretrained("fusing/latent-diffusion-text2im-large")
ldm_pipeline = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")

FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=milyiyo.testing-diffusers" />'

def greet(name):
    return "Hello " + name + "!!"
    
def genimage(prompt, iterations):
  image = pipeline([prompt], generator=generator, eta=0.3, guidance_scale=6.0, num_inference_steps=iterations)["sample"]
  return image[0]

  #image_processed = image.cpu().permute(0, 2, 3, 1)
  #image_processed = image_processed  * 255.
  #image_processed = image_processed.numpy().astype(np.uint8)
  #image_pil = PIL.Image.fromarray(image_processed[0])

  # save image
  #file_name = "test.png"
  #image_pil.save(file_name)
  #img = cv2.imread(file_name)
  ##cv2_imshow(img)
  #return img
  
def predict(prompt, steps=100):
  torch.cuda.empty_cache()
  generator = torch.manual_seed(42)
  images = ldm_pipeline([prompt], generator=generator, num_inference_steps=steps, eta=0.3, guidance_scale=6.0)["sample"]
  return images[0]

iface = gr.Interface(
  fn=predict, 
  inputs=["text", "number"], 
  outputs=gr.Image(shape=[256,256], type="pil", elem_id="output_image"))
iface.launch()