import gradio as gr
import numpy as np
from kandinsky2 import get_kandinsky2
import torch


def text2img(text,heig,widt,num,sampler,numsteps,batchsize):
    torch.cuda.empty_cache()
    images =[]
    model = get_kandinsky2('cuda', task_type='text2img', cache_dir='/tmp/kandinsky2', model_version='2.1', use_flash_attention=False)
    for i in range(num):
        outputs = model.generate_text2img(
                text, 
                num_steps=numsteps,
                batch_size=batchsize, 
                guidance_scale=4,
                h=heig, w=widt,
                sampler=sampler, 
                prior_cf_scale=4,
                prior_steps="5"
                    
                )
        images.append(outputs[0])
    return images
def mixing(image1,image2,sampler,num_steps,batch_size,num):
        torch.cuda.empty_cache()
        images =[]
        model = get_kandinsky2('cuda', task_type='text2img', model_version='2.1', use_flash_attention=False)
        for i in range(num):
            weights = [0.5, 0.5]
            output_mixing = model.mix_images(
                [image1,image2], 
                weights, 
                num_steps,
                batch_size, 
                guidance_scale=5,
                h=300, w=300,
                sampler=sampler, 
                prior_cf_scale=4,
                prior_steps="5"
            ) 
            images.append(output_mixing[0])
        return images
with gr.Blocks()as demo:
    gr.Markdown("Генерация изображения")
    with gr.Tab("Генерация изображения по текстовому запросу"):
        with gr.Blocks():
            with gr.Row().style(equal_height=True):
                    with gr.Column():
                        text_input=gr.Textbox()
                        sampler = gr.Radio(
                            ["ddim_sampler", "p_sampler", "plms_sampler"],
                            value="p_sampler",
                            label="Sampler",)
                        size_h=gr.Slider(216,864,step=100,label="Высота")
                        size_w=gr.Slider(216,864,step=100,label="Ширина")
                        num_steps=gr.Slider(50,250,step=25,label="Количество шагов")
                        batch_size=gr.Slider(2,20,step=1,label="Размер батча")
                        num=gr.Slider(1,8,step=1,label="Количество")        
                    with gr.Column():
                        image_output=gr.Gallery()
                        image_button = gr.Button("Сгенерировать")
                        image_button.click(text2img,inputs=[text_input,size_h,size_w,num,sampler,num_steps,batch_size],outputs=image_output)
    with gr.Tab("Объеденение двух изображений"):
        with gr.Row():
            with gr.Column():
                image_input1=gr.Image(type="pil")
                image_input2=gr.Image(type="pil")
                sampler = gr.Radio(
                            ["ddim_sampler", "p_sampler", "plms_sampler"],
                            value="p_sampler",
                            label="Sampler",)
                num_steps=gr.Slider(50,250,step=25,label="Количество шагов")
                batch_size=gr.Slider(2,20,step=1,label="Размер батча")
                num=gr.Slider(1,8,step=1,label="Количество")
            with gr.Column():
                image_mixing_output=gr.Gallery(type="pil")
                image_mixing_output_button=gr.Button("Сгенерировать")
        image_mixing_output_button.click(mixing,inputs=[image_input1,image_input2,sampler,num_steps,batch_size,num],outputs=image_mixing_output)
demo.launch(share=True)