import torch
import gradio as gr
import torch
import os
from PIL import Image
from torch import autocast
from perpneg_diffusion.perpneg_stable_diffusion.pipeline_perpneg_stable_diffusion import PerpStableDiffusionPipeline

has_cuda = torch.cuda.is_available()
device = torch.device('cpu' if not has_cuda else 'cuda')
print(device)

# initialize stable diffusion model
pipe = PerpStableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    # use_auth_token=True
).to(device)

def dummy(images, **kwargs):
    return images, False


pipe.safety_checker = dummy

examples = [
    [
        "an armchair in the shape of an avocado | cushion in the armchair",
        "1 | -0.3",
        "145",
        "7.5"
    ],
    [
        "an armchair in the shape of an avocado",
        "1",
        "145",
        "7.5"
    ],
    [
        "a peacock, back view | a peacock, front view",
        "1 | -3.5",
        "30",
        "7.5"
    ],
    [
        "a peacock, back view",
        "1",
        "30",
        "7.5"
    ],    
    [
        "A boy wearing sunglasses | a pair of sunglasses with white frame",
        "1 | -0.35",
        "200",
        "11"
    ],
    [
        "A boy wearing sunglasses",
        "1",
        "200",
        "11",
    ],
    [
        "a photo of an astronaut riding a horse | a jumping horse | a white horse", 
        "1 | -0.3 | -0.1",
        "1988",
        "10"
    ],             
    [
        "a photo of an astronaut riding a horse | a jumping horse", 
        "1 | -0.3",
        "1988",
        "10"
    ],
    [
        "a photo of an astronaut riding a horse", 
        "1",
        "1988",
        "10"
    ],         
]







def predict(prompt, weights, seed, scale=7.5, steps=50):
    try:
        with torch.no_grad():
            has_cuda = torch.cuda.is_available()
            with autocast('cpu' if not has_cuda else 'cuda'):
                if has_cuda:
                    generator = torch.Generator('cuda').manual_seed(int(seed))
                else:
                    generator = torch.Generator().manual_seed(int(seed))
                image_perpneg = pipe(prompt, guidance_scale=float(scale), generator=generator,
                            num_inference_steps=steps, weights=weights)["images"][0]
                return image_perpneg
    except Exception as e:
        print(e)
        return None



MESSAGE = '''
Our method helps you achieve three amazing things:

1. Edit your generated images iteratively without damaging any important concepts.
2. Generate any view of objects that the original Stable Diffusion implementation couldn't produce. For example, you can generate a "peacock, back view" by using "peacock, front view" as the negative prompt. Compare our method to [Stable Diffusion](https://huggingface.co/spaces/stabilityai/stable-diffusion).
3. Alleviate the multihead problem in text-to-3D. Check out our work on this at [perp-neg.github.io](https://perp-neg.github.io/).

To use our demo, simply enter your main prompt first, followed by a set of positive and negative prompts separated by "|". When only one prompt is provided and the weight of that prompt is 1, it is identical to using Stable Diffusion. We provided those as examples for the sake of comparison of our algorithm to Stable Diffusion. Put the weight of main prompt as 1. Provide a complete sentence for negative prompt. The number of weights should be equal to the number of the prompts. Vary the weight of the negative prompts from -0.1 to -3 to produce desired results.
Use our demo to create some amazing and unique images!
'''

MESSAGE_END = '''

Unlike the original implementation, our method ensures that everything provided as the main prompt remains intact even when there is an overlap between the positive and negative prompts.

We've also integrated the idea of robust view generation in text-to-3D to avoid the multihead problem. For more details, please check out our work on this at [perp-neg.github.io](https://perp-neg.github.io/).

'''

app = gr.Blocks()
with app:
    # gr.Markdown(
    #     "# **<p align='center'>AMLDS Video Tagging</p>**"
    # )
    gr.Markdown(
        "# **<p align='center'>Perp-Neg: Iterative Editing and Robust View Generation</p>**"
    )    
    gr.Markdown(
        """
        ### **<p align='center'>Demo created by Reza Armandpour and Huangjie Zheng</p>**
        """
    )
    gr.Markdown(MESSAGE)
 
    with gr.Row():
        with gr.Column():
            # with gr.Tab(label="Inputs"):
            # gr.Markdown(
            #     "### Prompts (a list of prompts separated by vertical bar | )"
            # )
            prompt = gr.Textbox(label="Prompts (a list of prompts separated by vertical bar | ):", show_label=True, placeholder="a peacock, back view | a peacock, front view")
            weights = gr.Textbox(
                label="Weights (a list of weights separated by vertical bar | )", show_label=True, placeholder="1 | -3.5"
            )
            seed = gr.Textbox(
                label="Seed", show_label=True, value=30
            )
            scale = gr.Textbox(
                label="Guidance scale", show_label=True, value=7.5
            )                   
            image_gen_btn = gr.Button(value="Generate")

        with gr.Column():
            img_output = gr.Image(
                label="Result",
                show_label=True,
            )


    gr.Markdown("**Examples:**")
    gr.Examples(
        examples,
        [prompt, weights, seed, scale],
        [img_output],
        fn=predict,
        cache_examples=False,
    )

    image_gen_btn.click(
        predict,
        inputs=[prompt, weights, seed, scale],
        outputs=[img_output],
    )
    gr.Markdown("""
    \n The algorithem is based on the paper: [Re-imagine the Negative Prompt Algorithm: Transform 2D Diffusion into 3D, alleviate Janus problem and Beyond.](https://Perp-Neg.github.io)
    """)
    gr.Markdown(MESSAGE_END)       

    gr.Markdown(
        """
        \n Demo created by: Reza Armandpour and Huangjie Zheng</a>.
        """
    )

app.launch()