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

accelerator = Accelerator(cpu=True)
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")
generator = torch.Generator(device="cpu").manual_seed(random.randint(1, 4876364))

def plex(img, cook, one, two, three):
    goof = load_image(img).resize((512, 512))
    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=5, generator=generator)
    negative_emb = pope_prior(prompt=negative_prior_prompt, guidance_scale=1, num_inference_steps=5, generator=generator)
    imags = pope(image_embeds=img_emb.image_embeds,negative_image_embeds=negative_emb.image_embeds,num_inference_steps=10,generator=generator,height=512,width=512,).images[0]
    images_texts = [cook, goof, imags]
    weights = [one, two, three]
    primpt = ""
    prior_out = pope_prior.interpolate(images_texts, weights, num_inference_steps=5)
    imas = pope(**prior_out, height=512, width=512, num_inference_steps=25).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.queue(max_size=1)
iface.launch(max_threads=1)