import gradio as gr
import numpy as np
import tensorflow as tf
from huggingface_hub import from_pretrained_keras
from diffusion_sampler import DiffusionSampler

devices = '\n'.join([f'- {device.name}' for device in tf.config.list_physical_devices('GPU')]) or 'No GPU devices found.'
print(f"GPUs available: {devices}")

scheduler_button = gr.Radio(
    choices=["Linear", "Cosine"],
    label="Noise Scheduler",
    value="Linear",
    info="""
        Decides whether to employ a model trained with a linear scheduler, 
        as proposed by Jonathan Ho et al. in 'Denoising Diffusion Probabilistic Models',
        or the cosine variant introduced by Alex Nichol et al. in 'Improved Denoising Diffusion Probabilistic Models'.
    """,
)

sampling_button = gr.Radio(
    choices=["DDPM", "DDIM"],
    label="Sampling Procedure",
    value="DDPM",
    info="""
        Selects either the stocasthic sampling procedure described by Jonathan Ho et al. in 'Denoising Diffusion Probabilistic Models',
        or the implicit variant proposed by Jiaming Song et al. in 'Denoising Diffusion Implicit Models'.
        For the latter, it is also necessary to specify the sub-sequence strategy and the number of sampling steps.
    """,
)

subsequence_button = gr.Radio(
    choices=["Linear", "Quadratic"],
    label="Sub-Sequence",
    value="Linear",
    info="""
        Specific to DDIM sampling, this parameter chooses the procedure 
        for forming the sub-sequence employed during the sampling process.
    """,
)

ema_button = gr.Checkbox(
    value=True,
    label="Exponential Moving Average",
    info="""
        Whether to invoke the network with the applied exponential moving average on the model's weights.
    """
)

images_button = gr.Number(
    label="Number of images to generate",
    value=10,
    precision=0,
    minimum=1,
    maximum=64,
    info="""
        The number of images to be generated.
        Larger batch sizes result in longer inference times.
    """
)

step_button = gr.Slider(
    minimum=700,
    value=1_000,
    maximum=1_000,
    randomize=True,
    label="Number of sampling steps",
    info="""
        Relevant exclusively to DDIM sampling, this parameter determines the number of steps to be utilized during sampling.
        The default value is set to 1000 in the case of DDPM sampling.
    """
)

gallery = gr.Gallery(
    columns=4,
    allow_preview=False,
    show_download_button=False,
    show_share_button=False,
    label="""
        Generated Flowers
    """
)

diffusion_models = {
    "linear":
        DiffusionSampler(
            model=from_pretrained_keras("leowajda/linear_diffusion", cache_dir="cache"),
            ema_model=from_pretrained_keras("leowajda/linear_diffusion_ema", cache_dir="cache"),
            noise_scheduler="linear"
        ),

    "cosine":
        DiffusionSampler(
            model=from_pretrained_keras("leowajda/cosine_diffusion", cache_dir="cache"),
            ema_model=from_pretrained_keras("leowajda/cosine_diffusion_ema", cache_dir="cache"),
            noise_scheduler="cosine"
        )
}


def call_model(
        model_to_call: str,
        sample_strategy: str = "ddim",
        step_strategy: str = "uniform",
        ema: bool = True,
        steps: int = 1_000,
        num_images: int = 0,
):
    diffusion_model = diffusion_models[model_to_call.lower()]
    images = diffusion_model.generate_images(
        num_images=num_images,
        steps=steps,
        sample_strategy=sample_strategy.lower(),
        step_strategy=step_strategy.lower(),
        ema=ema,
    )

    return images.numpy().astype(np.uint8)


demo = gr.Interface(
    fn=call_model,
    inputs=[scheduler_button, sampling_button, subsequence_button, ema_button, step_button, images_button],
    outputs=gallery,
    cache_examples=True,
    title="""Unconditional Image Generation Through Denoising Diffusion Implicit Models""",
    examples=[
        ["Linear", "DDPM", "Linear", True, 1_000, 25],
        ["Linear", "DDPM", "Linear", False, 1_000, 25],
        ["Cosine", "DDPM", "Linear", False, 750, 25],
        ["Cosine", "DDPM", "Linear", True, 750, 25],
        ["Linear", "DDIM", "Linear", True, 750, 25],
    ],
    description="""                
        <p align="center">
            Supervisor: <strong>Wojciech Oronowicz – Jaśkowiak, PhD</strong>
            &emsp;
            Author: <strong>Leonardo Wajda</strong>
            &emsp;
            Specialization: <strong>Intelligent Data Processing Systems</strong>
        </p>
    """,
)

demo.queue(default_concurrency_limit=None).launch()