import os
import gradio as gr
import torch

from diffusers import StableDiffusionImg2ImgPipeline

os.environ['GRADIO_THEME'] = 'default'
# load the pipeline
device = "cpu"
# model_id_or_path = "runwayml/stable-diffusion-v1-5"
# takes too much time downloading, change to local path
model_path = "."
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_path)
img2img_pipe = img2img_pipe.to(device)


def img2img_diff(prompt, img):
    return img2img_pipe(prompt=prompt, image=img, strength=0.75, guidance_scale=7.5).images[0]


gr.Interface(fn=img2img_diff, inputs=['text', 'image'], outputs='image')