import gradio as gr
import torch
import numpy as np
from PIL import Image
from datasets import load_dataset
from diffusers import StableDiffusionImg2ImgPipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2")
pipe = pipe.to(device)

def resize(value,img):
    img = Image.open(img)
    img = img.resize((value,value))
    return img

def infer(source_img, prompt, guide, steps, seed, Strength):
    generator = torch.Generator(device).manual_seed(seed)     
    source_image = resize(768, source_img)
    source_image.save('source.png')
    image = pipe([prompt], init_image=source_image, strength=Strength, guidance_scale=guide, num_inference_steps=steps).images[0]
    return image

gr.Interface(fn=infer, inputs=[gr.Image(source="upload", type="filepath", label="Raw Image"), gr.Textbox(label = 'Prompt Input Text'),
    gr.Slider(2, 15, value = 7, label = 'Guidence Scale'),
    gr.Slider(10, 50, value = 25, step = 1, label = 'Number of Iterations'),
    gr.Slider(
        label = "Seed",
        minimum = 0,
        maximum = 2147483647,
        step = 1,
        randomize = True), gr.Slider(label='Strength', minimum = 0, maximum = 1, step = .05, value = .5)
    ], outputs='image', title = "Stable Diffusion 2.0 Image to Image Pipeline CPU", description = "For more information on Stable Diffusion 2.0 see https://github.com/Stability-AI/stablediffusion <br><br>Upload an Image (must be .PNG and 512x512-2048x2048) enter a Prompt, or let it just do its Thing, then click submit. 10 Iterations takes about 900 seconds currently. For more informationon about Stable Diffusion or Suggestions for prompts, keywords, artists or styles see https://github.com/Maks-s/sd-akashic", article = "Code Monkey: <a href=\"https://huggingface.co/Manjushri\">Manjushri</a>").queue(max_size=10).launch(debug=True)