import gradio as gr
import os
from PIL import Image, ImageOps
import matplotlib.pyplot as plt
import numpy as np
import torch
import requests
from tqdm import tqdm

from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline
import torchvision.transforms as T

from utils import preprocess,prepare_mask_and_masked_image, recover_image

to_pil = T.ToPILImage()

model_id_or_path = "runwayml/stable-diffusion-v1-5"
# model_id_or_path = "CompVis/stable-diffusion-v1-4"
# model_id_or_path = "CompVis/stable-diffusion-v1-3"
# model_id_or_path = "CompVis/stable-diffusion-v1-2"
# model_id_or_path = "CompVis/stable-diffusion-v1-1"

pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id_or_path,
    revision="fp16", 
    torch_dtype=torch.float16,
)
pipe_img2img = pipe_img2img.to("cuda")

pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    revision="fp16",
    torch_dtype=torch.float16,
)
pipe_inpaint = pipe_inpaint.to("cuda")

def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1, mask=None):
    X_adv = X.clone().detach() + (torch.rand(*X.shape)*2*eps-eps).cuda()
    pbar = tqdm(range(iters))
    for i in pbar:
        actual_step_size = step_size - (step_size - step_size / 100) / iters * i  

        X_adv.requires_grad_(True)

        loss = (model(X_adv).latent_dist.mean).norm()

        pbar.set_description(f"[Running attack]: Loss {loss.item():.5f} | step size: {actual_step_size:.4}")

        grad, = torch.autograd.grad(loss, [X_adv])
        
        X_adv = X_adv - grad.detach().sign() * actual_step_size
        X_adv = torch.minimum(torch.maximum(X_adv, X - eps), X + eps)
        X_adv.data = torch.clamp(X_adv, min=clamp_min, max=clamp_max)
        X_adv.grad = None    
        
        if mask is not None:
            X_adv.data *= mask
            
    return X_adv
    
def process_image_img2img(raw_image,prompt):
    resize = T.transforms.Resize(512)
    center_crop = T.transforms.CenterCrop(512)
    init_image = center_crop(resize(raw_image))
    with torch.autocast('cuda'):
        X = preprocess(init_image).half().cuda()
        adv_X = pgd(X, 
                    model=pipe_img2img.vae.encode, 
                    clamp_min=-1, 
                    clamp_max=1,
                    eps=0.06, # The higher, the less imperceptible the attack is 
                    step_size=0.02, # Set smaller than eps
                    iters=100, # The higher, the stronger your attack will be
                   )
        
        # convert pixels back to [0,1] range
        adv_X = (adv_X / 2 + 0.5).clamp(0, 1)

    adv_image = to_pil(adv_X[0]).convert("RGB")
    
    # a good seed (uncomment the line below to generate new images)
    SEED = 9222
    # SEED = np.random.randint(low=0, high=10000)
    
    # Play with these for improving generated image quality
    STRENGTH = 0.5
    GUIDANCE = 7.5
    NUM_STEPS = 50
    
    with torch.autocast('cuda'):
        torch.manual_seed(SEED)
        image_nat = pipe_img2img(prompt=prompt, image=init_image, strength=STRENGTH, guidance_scale=GUIDANCE, num_inference_steps=NUM_STEPS).images[0]
        torch.manual_seed(SEED)
        image_adv = pipe_img2img(prompt=prompt, image=adv_image, strength=STRENGTH, guidance_scale=GUIDANCE, num_inference_steps=NUM_STEPS).images[0]

    return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]

def process_image_inpaint(raw_image,mask, prompt):
    init_image = raw_image.convert('RGB').resize((512,512))
    mask_image = mask.convert('RGB')
    mask_image = ImageOps.invert(mask_image).resize((512,512))

    # Attack using embedding of random image from internet
    target_url = "https://bostonglobe-prod.cdn.arcpublishing.com/resizer/2-ZvyQ3aRNl_VNo7ja51BM5-Kpk=/960x0/cloudfront-us-east-1.images.arcpublishing.com/bostonglobe/CZOXE32LQQX5UNAB42AOA3SUY4.jpg"
    response = requests.get(target_url)
    target_image = Image.open(BytesIO(response.content)).convert("RGB")
    target_image = target_image.resize((512, 512))

    with torch.autocast('cuda'):
    mask, X = prepare_mask_and_masked_image(init_image, mask_image)
    X = X.half().cuda()
    mask = mask.half().cuda()
    
    # Here we attack towards the embedding of a random target image. You can also simply attack towards an embedding of zeros!
    target = pipe_inpaint.vae.encode(preprocess(target_image).half().cuda()).latent_dist.mean

    adv_X = pgd(X, 
                target = target,
                model=pipe_inpaint.vae.encode, 
                criterion=torch.nn.MSELoss(), 
                clamp_min=-1, 
                clamp_max=1,
                eps=0.06, 
                step_size=0.01, 
                iters=1000,
                mask=1-mask
               )

    adv_X = (adv_X / 2 + 0.5).clamp(0, 1)

    adv_image = to_pil(adv_X[0]).convert("RGB")
    adv_image = recover_image(adv_image, init_image, mask_image, background=True)

    # A good seed
    SEED = 9209
    
    # Uncomment the below to generated other images
    # SEED = np.random.randint(low=0, high=100000)
    
    torch.manual_seed(SEED)
    print(SEED)
    
    strength = 0.7
    guidance_scale = 7.5
    num_inference_steps = 100
    
    image_nat = pipe_inpaint(prompt=prompt, 
                         image=init_image, 
                         mask_image=mask_image, 
                         eta=1,
                         num_inference_steps=num_inference_steps,
                         guidance_scale=guidance_scale,
                         strength=strength
                        ).images[0]
    image_nat = recover_image(image_nat, init_image, mask_image)
    
    torch.manual_seed(SEED)
    image_adv = pipe_inpaint(prompt=prompt, 
                         image=adv_image, 
                         mask_image=mask_image, 
                         eta=1,
                         num_inference_steps=num_inference_steps,
                         guidance_scale=guidance_scale,
                         strength=strength
                        ).images[0]
    image_adv = recover_image(image_adv, init_image, mask_image)

    return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]


examples = [["dog.png", "dog under heavy rain and muddy ground real"]]


with gr.Blocks() as demo:
    gr.Markdown("""
                ## Interactive demo: Raising the Cost of Malicious AI-Powered Image Editing
            """)
    gr.HTML('''
     <p style="margin-bottom: 10px; font-size: 94%">This is an unofficial demo for Photoguard, which is an approach to safe-guarding images against manipulation by ML-powerd photo-editing models such as stable diffusion through immunization of images. The demo is based on the <a href='https://github.com/MadryLab/photoguard' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors.</p>
              ''')
    
    with gr.Column(): 
        with gr.Tab("Simple Image to Image"):
            input_image_img2img = gr.Image(type="pil", label = "Source Image")
            input_prompt_img2img = gr.Textbox(label="Prompt")
            run_btn_img2img = gr.Button('Run')
        
        with gr.Tab("Simple Inpainting"):
            input_image_inpaint = gr.Image(type="pil", label = "Source Image")
            mask_image_inpaint = gr.Image(type="pil", label = "Mask")
            input_prompt_inpaint = gr.Textbox(label="Prompt")
            run_btn_inpaint = gr.Button('Run')
        
        with gr.Row():
            result_gallery = gr.Gallery(
            label="Generated images", show_label=False, elem_id="gallery"
        ).style(grid=[2], height="auto")
            
    run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img], outputs = [result_gallery])
    run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
    
                     
demo.launch(debug=True)