import gradio as gr
import torch
from PIL import Image

from lambda_diffusers import StableDiffusionImageEmbedPipeline

def main(
    input_im,
    scale=3.0,
    n_samples=4,
    steps=25,
    seed=0,
    ):
    
    generator = torch.Generator(device=device).manual_seed(int(seed))

    images_list = pipe(
        n_samples*[input_im],
        guidance_scale=scale,
        num_inference_steps=steps,
        generator=generator,
        )

    images = []
    for i, image in enumerate(images_list["sample"]):
        if(images_list["nsfw_content_detected"][i]):
            safe_image = Image.open(r"unsafe.png")
            images.append(safe_image)
        else:
            images.append(image)
    return images

device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionImageEmbedPipeline.from_pretrained(
    "lambdalabs/sd-image-variations-diffusers",
    revision="273115e88df42350019ef4d628265b8c29ef4af5",
    )
pipe = pipe.to(device)

inputs = [
    gr.Image(),
    gr.Slider(0, 25, value=3, step=1, label="Guidance scale"),
    gr.Slider(1, 4, value=1, step=1, label="Number images"),
    gr.Slider(5, 50, value=25, step=5, label="Steps"),
    gr.Number(0, labal="Seed", precision=0)
]
output = gr.Gallery(label="Generated variations")
output.style(grid=2)

demo = gr.Interface(
    fn=main,
    title="Stable Diffusion Image Variations",
    inputs=inputs,
    outputs=output
    )
demo.launch()