import gradio as gr
import inspect
import warnings
import numpy as np
from typing import List, Optional, Union
import requests

from io import BytesIO
from PIL import Image
import torch
from torch import autocast
from tqdm.auto import tqdm
from diffusers import StableDiffusionImg2ImgPipeline

from huggingface_hub import notebook_login
notebook_login()


device = "cuda"
model_path = "CompVis/stable-diffusion-v1-4"
access_token = "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"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_path,
    revision="fp16", 
    torch_dtype=torch.float16,
    use_auth_token=access_token
)
pipe = pipe.to(device)

def predict(img, strength, seed, prompt):
  seed =  int(seed)
  img1 = np.asarray(img)
  img2 = Image.fromarray(img1)

  init_image = img2.resize((768, 512))

  generator = torch.Generator(device=device).manual_seed(seed)
  with autocast("cuda"):
    image = pipe(prompt=prompt, init_image=init_image, strength=strength, guidance_scale=5, generator=generator).images[0]

  return image

gr.Interface(
    predict,
    title = 'Image to Image using Diffusers',
    inputs=[
        gr.Image(),
        gr.Slider(0, 1, value=0.05, label ="strength (keep it close to 0 to make minimal changes to image (such as 0.1, 0.2, 0.3)"),
        gr.Number(label = "seed (any number, generally 1024. But it's totally random. Change it and see different outputs)"),
        gr.Textbox(label="Prompt, empty by default")
    ],
    outputs = [
        gr.Image()
        ]
).launch()