import gradio as gr
import torch
# region_offset = torch.tensor(region_offset).int()

from utils import gen_image_as_per_prompt

styles = ["depthmap", "cosmicgalaxy", "concept-art", "Marc Allante", "midjourney-style", "No style"]
styleValues = ["learned_embeds_depthmap.bin",
               "learned_embeds_cosmic-galaxy-characters-style.bin",
               "learned_embeds_sd_concept-art.bin",
               "learned_embeds_style-of-marc-allante.bin",
               "learned_embeds_midjourney.bin",
               ""]
seed_values = [30, 24, 35, 47, 78, 42]

styles_dict = dict(zip(styles, styleValues))
seed_dict = dict(zip(styles, seed_values))


# Custom loss function
def reduce_highlight(images):
    """Calculates the mean absolute error for amber color.

    Args:
      images: A tensor of shape (batch_size, channels, height, width).
      target_red: Target red value for amber.
      target_green: Target green value for amber.
      target_blue: Target blue value for amber.

    Returns:
      The mean absolute error.
      #target_red=0.8, target_green=0.6, target_blue=0.4
    """

    red_error = torch.abs(images[:, 0] - 0.12).mean()
    green_error = torch.abs(images[:, 1] - 0.2).mean()
    blue_error = torch.abs(images[:, 2] - 0.15).mean()

    # You can adjust weights for each channel if needed
    amber_error = (red_error + green_error + blue_error) / 3
    return amber_error

def _inference(text, style, use_loss=False):
    if use_loss:
        image = gen_image_as_per_prompt(text, styles_dict[style], seed_dict[style], reduce_highlight)
    else:
        image = gen_image_as_per_prompt(text, styles_dict[style], seed_dict[style])
    return image


title = "Stable Diffusion with different styles"
description = "In this demo, the word 'puppy' is replaced by the style that is selected"
examples = [["oil painting of a dragon in puppy style", "mosiac", True],
            ["Spiderman in puppy style", "midjourney", True],
            ["Batman in puppy style", "matrix", False],
             ["Mojo Jojo in puppy style", "No style", False]]

demo = gr.Interface(
    _inference,
    inputs=[
        gr.Textbox(placeholder="Type a prompt with word 'puppy' in it..", container=False, scale=7),
        gr.Radio(styles, label="Select a Style"),
        gr.Checkbox(label="Use custom loss")
    ],
    outputs=[
        gr.Image(width=256, height=256, label="output")
        # gr.Text(label="output")
    ],
    title=title,
    description=description,
    examples=examples,
    cache_examples=False
)
demo.launch(debug=True)