import os
import gradio as gr
import torch

print(f"Is CUDA available: {torch.cuda.is_available()}")
# True
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")

from diffusers import StableDiffusionImg2ImgPipeline

os.environ['GRADIO_THEME'] = 'default'
# load the pipeline
device = "cuda"
model_id_or_path = "runwayml/stable-diffusion-v1-5"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
img2img_pipe = img2img_pipe.to(device)


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


app = gr.Interface(fn=img2img_diff, inputs=[gr.Text(label="prompt text"), gr.Image(type='pil', label='draft image')],
                   outputs=gr.Image(type='pil'))

app.launch(debug=True)