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

access_token = "TOKEN"
# Go to hugging face, your profile "SETTINGS" options, click on "Access tokens", and then generate ur token from there. Paste it over the top

# load the pipeline
device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    revision="fp16", 
    torch_dtype=torch.float16,
use_auth_token=access_token
).to(device)






def generate(img, strength, seed, prompt):
    # Convert the seed to an integer
    seed = int(seed)

    img1 = np.asarray(img)
    img2 = Image.fromarray(img1)

    # Check that the input image is a valid image
    if not isinstance(img2, Image.Image):
        raise ValueError("Invalid input image")

    # Resize the image
    init_image = img2.resize((768, 512))

    # Create a list to store the 4 output images
    images = []

    # Use the GPU if available, otherwise use the CPU
    device = "cuda" if torch.cuda.is_available() else "cpu"
    with autocast(device):
        # Generate the 4 output images using the pipe function
        for i in range(2):
            # Initialize the generator with a random seed
            generator = torch.Generator(device=device).manual_seed(seed*i)

            # Call the pipe function and store the output image
            output_image = pipe(prompt=prompt, init_image=init_image, strength=strength, guidance_scale=7.5, generator=generator, batch_size=128).images[0]

            # Check that the output image is a valid image
            if not isinstance(output_image, Image.Image):
                raise ValueError("Invalid output image")

            images.append(output_image)

    return [images[0], images[1]]


gr.Interface(
    
    generate,
    title = 'Image to Image using Diffusers',
    inputs=[
        gr.Image(elem_id = "input-image"),
        gr.Slider(0, 1, value=0.05, label ="Strength (keep close to 0 for minimal changes)"),
        gr.Slider(50, 700, value=75, label ="Seed"),
        gr.Textbox(label="Prompt (leave blank if you want minimal changes)"),
    ],
    outputs = [
        gr.Image(elem_id="output-image"),
        gr.Image(elem_id="output-image"),

        ], css = "#output-image, #input-image, #image-preview {border-radius: 40px !important; background-color : gray !important;} "
).launch(share=True, debug=True)