import torch
import kornia
import gradio as gr


def image_to_grayscale(image, version, step_no):
    image = kornia.utils.image_to_tensor(image, keepdim=False) / 255.
    image = kornia.geometry.transform.resize(
        image, 384, interpolation='bilinear', align_corners=None, side='short', antialias=False)
    dissolved = kornia.filters.StableDiffusionDissolving(version)(image, int(step_no))
    output = kornia.utils.tensor_to_image(dissolved)
    return output


iface = gr.Interface(
    fn=image_to_grayscale,
    inputs=[
        gr.Image(type="numpy"),
        gr.Dropdown(
            ["2.1", "1.5", "1.4"], value="2.1", multiselect=False, label="Stable Diffusion Version"
        ),
        gr.Number(value=500, minimum=0, precision=0, maximum=1000, label="Timestep No."),
    ],
    outputs="image",
    title="Dissolving Transformations with Stable Diffusion",
    description="Dissolving transformation based on `Dissolving Is Amplifying: Towards Fine-Grained Anomaly Detection`. Original rep is https://github.com/shijianjian/DIA. It has been included in Kornia lib (https://kornia.readthedocs.io/en/latest/augmentation.module.html#kornia.augmentation.RandomDissolving). You may try by `kornia.filters.StableDiffusionDissolving(version)(image, step_no)`, or `kornia.augmentation.RandomDissolving()(image)`. In this demo, the images are downsampled for fast computation."
)

# Launch the Gradio app
iface.launch()