import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL
from PIL import Image
from ip_adapter.ip_adapter_faceid import IPAdapterFaceIDPlus
import cv2
from insightface.app import FaceAnalysis
from insightface.utils import face_align
import gradio as gr
from huggingface_hub import hf_hub_download
from datetime import datetime


def download_models():
    hf_hub_download(
        repo_id='h94/IP-Adapter-FaceID',
        filename='ip-adapter-faceid-plus_sd15.bin',
        local_dir='IP-Adapter-FaceID')
    hf_hub_download(
        repo_id='h94/IP-Adapter',
        filename='models/image_encoder/config.json',
        local_dir='IP-Adapter')
    hf_hub_download(
        repo_id='h94/IP-Adapter',
        filename='models/image_encoder/pytorch_model.bin',
        local_dir='IP-Adapter')


def get_ip_model():
    download_models()
    base_model_path = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
    vae_model_path = "stabilityai/sd-vae-ft-mse"
    image_encoder_path = "IP-Adapter/models/image_encoder"
    ip_ckpt = "IP-Adapter-FaceID/ip-adapter-faceid-plus_sd15.bin"

    if torch.cuda.is_available():
        device = 'cuda'
        torch_dtype = torch.float16
    else:
        device = 'cpu'
        torch_dtype = torch.float32
    print(f'Using device: {device}')

    noise_scheduler = DDIMScheduler(
        num_train_timesteps=1000,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="scaled_linear",
        clip_sample=False,
        set_alpha_to_one=False,
        steps_offset=1,
    )
    vae = AutoencoderKL.from_pretrained(vae_model_path).to(dtype=torch_dtype)
    pipe = StableDiffusionPipeline.from_pretrained(
        base_model_path,
        torch_dtype=torch_dtype,
        scheduler=noise_scheduler,
        vae=vae,
        feature_extractor=None,
        safety_checker=None
    )

    ip_model = IPAdapterFaceIDPlus(pipe, image_encoder_path, ip_ckpt, device, num_tokens=4, torch_dtype=torch_dtype)
    return ip_model


ip_model = get_ip_model()
app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.2)

def generate_images(prompt, img_filepath,
                    negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality, blurry",
                    img_prompt_scale=0.5,
                    num_inference_steps=40,
                    seed=None, n_images=4):
    print(f'{datetime.now().strftime("%Y/%m/%d %H:%M:%S")}: {prompt}')
    image = cv2.imread(img_filepath)
    faces = app.get(image)

    faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0)
    face_image = face_align.norm_crop(image, landmark=faces[0].kps, image_size=224) # you can also segment the face
    images = ip_model.generate(
        prompt=prompt, negative_prompt=negative_prompt, face_image=face_image, faceid_embeds=faceid_embeds,
        num_samples=n_images, width=512, height=512, num_inference_steps=num_inference_steps, seed=seed,
        scale=img_prompt_scale, # with scale=1 I get weird images
    )
    return [images[0], Image.fromarray(face_image[..., [2, 1, 0]])]


with gr.Blocks(css=".title { color: #4CAF50; font-family: 'Trebuchet MS', Helvetica, sans-serif; }") as demo:
    gr.Markdown(
        """
        <div class="title">
            # IP-Adapter-FaceID-plus
        </div>
        Generate images conditioned on an image prompt and a text prompt. Learn more here: [IP-Adapter-FaceID](https://huggingface.co/h94/IP-Adapter-FaceID).

        This demo is intended to use on GPU. It will work also on CPU but generating one image could take 900 seconds compared to a few seconds on GPU.
        """)
    
    with gr.Row():
        with gr.Column():
            demo_inputs = []
            demo_inputs.append(gr.Textbox(label='Text Prompt', value='Linkedin profile picture', placeholder="Enter a description for the image you want to generate", info="e.g., A professional profile picture, A singer on stage"))
            demo_inputs.append(gr.Image(type='filepath', label='Image Prompt', source="upload", tool="editor", info="Upload an image to guide the generation"))
            
            with gr.Accordion(label='Advanced Options', open=False):
                demo_inputs.append(gr.Textbox(label='Negative Text Prompt', value="monochrome, lowres, bad anatomy, worst quality, low quality, blurry", placeholder="Enter aspects to avoid in the generated image", info="e.g., blurry, low quality"))
                demo_inputs.append(gr.Slider(maximum=1, minimum=0, value=0.5, step=0.05, label='Image Prompt Scale', info="Adjust the influence of the image prompt"))
            
            btn = gr.Button("Generate")

        with gr.Column():
            demo_outputs = []
            demo_outputs.append(gr.Image(label='Generated Image', elem_id="generated-image"))
            demo_outputs.append(gr.Image(label='Detected Face', height=224, width=224, elem_id="detected-face"))

    btn.click(generate_images, inputs=demo_inputs, outputs=demo_outputs)
    
    sample_prompts = [
        'Linkedin profile picture',
        'A singer on stage',
        'A politician talking to the people',
        'An astronaut in space',
    ]
    gr.Examples(sample_prompts, inputs=demo_inputs[0], label='Sample Prompts')

demo.launch(share=True, debug=True)