from PIL import Image
import PIL.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=5,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=5).images[0]
    return imas

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