from PIL import Image
import cv2
import gradio as gr
import numpy as np
import torch
from accelerate import Accelerator
from transformers import pipeline
from diffusers.utils import load_image
from diffusers import KandinskyV22PriorPipeline, KandinskyV22Pipeline

accelerator = Accelerator(cpu=True)

generator = torch.Generator(device="cpu").manual_seed(4096)
pope_prior = accelerator.prepare(KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float32))
pope_prior = pope_prior.to("cpu")
    
pope = accelerator.prepare(KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float32))
pope = pope.to("cpu")

def plex(img, cook, one, two, three):
    goof = load_image(img).resize((512, 512))
    # We pass the prompt and negative prompt through the prior to generate image embeddings
    prompt = cook
    negative_prior_prompt = "lowres,text,bad quality,low quality,jpeg artifacts,ugly,bad hands,bad face,blurry,bad eyes,watermark,signature"
    img_emb = pope_prior(prompt=prompt, guidance_scale=0.85, num_inference_steps=10, generator=generator)
    negative_emb = pope_prior(prompt=negative_prior_prompt, guidance_scale=1, num_inference_steps=10, generator=generator)

    # run text2img pipeline
    imags = pope(
    image_embeds=img_emb.image_embeds,
    negative_image_embeds=negative_emb.image_embeds,
    num_inference_steps=20,
    generator=generator,
    height=512,
    width=512,
    ).images[0]

    ## return imags
    images_texts = [cook, goof, imags]

    # specify the weights for each condition in images_texts
    weights = [one, two, three]

    # We can leave the prompt empty
    primpt = ""
    prior_out = pope_prior.interpolate(images_texts, weights)
    imas = pope(**prior_out, height=512, width=512).images[0]
    return imas

iface = gr.Interface(fn=plex,inputs=[gr.Image(label="drop", type="pil"), gr.Textbox(label="prompt"), gr.Slider(label="Text Guide",minimum=0.01,step=0.01,maximum=1,value=0.5), gr.Slider(label="Your Image Guide",minimum=0.01,step=0.01,maximum=1,value=0.5),gr.Slider(label="Generated Image Guide",minimum=0.01,step=0.01,maximum=1,value=0.3)], outputs=gr.Image(), title="Ksky22 Cntrl Gdd Interp", description="ksky22 Cntrl Gdd Interp")
iface.launch()