import os
import gradio as gr
import torch
from torch import autocast
from kandinsky2 import get_kandinsky2
import multiprocessing as mp

device = torch.device('cpu')

from kandinsky2 import get_kandinsky2
model = get_kandinsky2('cpu', task_type='text2img', model_version='2.1', use_flash_attention=False)


def infer(args):
    prompt, negative = args
    with torch.no_grad():
        images = model.generate_text2img(prompt, 
                           negative_prior_prompt=negative,
                           negative_decoder_prompt=negative, 
                           num_steps=50,
                           batch_size=1,
                           guidance_scale=4,
                           h=768, w=768,
                           sampler='ddim_sampler',
                           prior_cf_scale=1,
                           prior_steps="25",)
    return images

def run_inference(prompt, negative):
    args = [(prompt, negative) for i in range(2)]
    with mp.Pool(2) as p:
        results = p.map(infer, args)
    return results[0]

css = """
        .gradio-container {
            font-family: 'IBM Plex Sans', sans-serif;
        }
        .gr-button {
            color: white;
            border-color: black;
            background: black;
        }
        input[type='range'] {
            accent-color: black;
        }
        .dark input[type='range'] {
            accent-color: #dfdfdf;
        }
        .container {
            max-width: 730px;
            margin: auto;
            padding-top: 1.5rem;
        }
        #gallery {
            min-height: 22rem;
            margin-bottom: 15px;
            margin-left: auto;
            margin-right: auto;
            border-bottom-right-radius: .5rem !important;
            border-bottom-left-radius: .5rem !important;
        }
        #gallery>div>.h-full {
            min-height: 20rem;
        }
        .details:hover {
            text-decoration: underline;
        }
        .gr-button {
            white-space: nowrap;
        }
        .gr-button:focus {
            border-color: rgb(147 197 253 / var(--tw-border-opacity));
            outline: none;
            box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);
            --tw-border-opacity: 1;
            --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);
            --tw-ring-shadow: var(--tw-ring-inset) 0 0