import gradio as gr
import inspect
import warnings
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"

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()