import gradio as gr
import argparse


import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from accelerate import Accelerator
from accelerate.utils import set_seed

import numpy as np


if __name__ == "__main__":
    accelerator = Accelerator()
    device = accelerator.device
    model_id = "stabilityai/stable-diffusion-2-1"
    
    # # SD args
    # resolution = "512,512"
    # num_inference_steps = 50
    # guidance_scale = 200
    # seed = 10
    # use_multidiffusion = False
    # use_skip_residual = False
    # use_dilated_sampling = False
    # use_progressive_upscaling = False
    # shuffle = False
    # use_md_prompt = False

    # load SD model
    pipe_pretrained = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe_pretrained.scheduler = DPMSolverMultistepScheduler.from_config(pipe_pretrained.scheduler.config)
    pipe_pretrained = pipe_pretrained.to(device)


    def generate_image(prompt,seed):
        set_seed(seed)
        image = pipe_pretrained(prompt,
                                width=128,
                                height=128, 
                                num_inference_steps=50
                                )[0]
        return image
    
    css = """
    #col-container{
        max-width: 720px;
        margin: 0 auto;
    }
    """

    MAX_SEED = np.iinfo(np.int32).max
    
    with gr.Blocks(css=css) as demo:
        with gr.Column(elem_id="col-container"):
            # Title and abstract
            gr.Markdown("# Single Layer Unlearning Gradient")
            gr.Markdown('<b><span style="color:red"> Notice: this demo is currently under debugging, complete features will be coming soon.</span></b>')
            gr.Markdown("### Abstract")
            gr.Markdown("""
            The unauthorized generation of privacy-related and copyright-infringing 
            content using generative-AI is becoming a significant concern for society, 
            raising ethical, legal, and privacy issues that demand urgent attention. 
            Recently, machine unlearning techniques have arisen that attempt to eliminate 
            the influence of sensitive content used during model training, but they often 
            require extensive updates in the model, reduce the utility of the models for unrelated content, 
            and/or incur substantial computational costs. 
            In this work, we propose a novel andefficient method called 
            Single Layer Unlearning Gradient (SLUG), that can unlearn targeted information by updating 
            a single targeted layer of a model using a one-time gradient computation. 
            We introduce two metrics: layer importance and gradient alignment, to identify 
            the appropriate layers for unlearning targeted information. 
            Our method is highly modular and enables selective removal of multiple concepts from the 
            generated outputs of widely used foundation models (e.g., CLIP), generative models (e.g., Stable Diffusion) and Vision-Language models.
            Our method shows effectiveness on a broad spectrum of concepts ranging from concrete 
            (e.g., celebrity name, intellectual property figure, and object) to abstract (e.g., novel concept and artistic style).
            """)
            # Buttons link to arxiv and github
            gr.HTML("""
            <div style="display:flex;column-gap:4px;">
                <a href='https://github.com/zikuicai/muwa'>
                    <img src='https://img.shields.io/badge/Code-github-blue'>
                </a>
                <a href='https://arxiv.org/abs/2407.11867'>
                    <img src='https://img.shields.io/badge/Paper-Arxiv-red'>
                </a>
                
            </div>
            """)
            # Interactive UI
            with gr.Group():
                prompt = gr.Textbox(label="Prompt", interactive=True)
                with gr.Row():
                    lr = gr.Slider(label="Unlearning step-size", minimum=0, maximum=5.00, step=0.2, value=0.0, interactive=True)
                    seed = gr.Slider(label="Random Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
                submit_btn = gr.Button("Submit", scale=2)
            
            gr.Examples(
                label="Example prompts",
                examples = [
                    ["A portrait of Elon Musk."],
                    ["A portrait of Mark Zuckerberg."],
                    ["Iron Man."],
                    ["Mickey Mouse."],
                ],
                inputs = [prompt]
            )
            
            output_images = gr.Gallery(label="Output Image", format="png", interactive=False)
        
        # Event Listeners
        submit_btn.click(
            fn = generate_image, 
            inputs = [prompt,seed],
            outputs = output_images,
            show_api=False
        )
        
    demo.launch(show_api=False, show_error=True)