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

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"

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

def predict(image_url, strength, seed):
  seed=  int(seed)
  
  response = requests.get(image_url)
  init_img = Image.open(BytesIO(response.content)).convert("RGB")
  init_img = init_img.resize((768, 512))


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

  return image
  
  
  
  gr.Interface(
    predict,
    title = 'Image to Image using Diffusers',
    inputs=[
        gr.Textbox(label="image_url"),

        gr.Slider(0, 1, value=0.05, label ="strength"),
        gr.Number(label = "seed")
        
    ],
    outputs = [
        gr.Image()
        ]
).launch()