Update app.py

app.py

@@ -1,20 +1,28 @@
 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' if torch.cuda.is_available() else 'cpu')
 
-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, 
+"""
+num_steps=50, 
+    batch_size=4, 
+    guidance_scale=7,
+    h=768, 
+    w=768,
+    sampler='ddim_sampler', 
+    prior_cf_scale=1,
+    prior_steps='25',
+"""
+def infer(prompt, negative='low quality, bad quality'):
+    images = model.generate_text2img(prompt, 
                            negative_prior_prompt=negative,
                            negative_decoder_prompt=negative, 
                            num_steps=50,
@@ -26,12 +34,6 @@ def infer(args):
                            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;
@@ -75,4 +77,176 @@ css = """
             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
+            --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);
+            --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));
+            --tw-ring-opacity: .5;
+        }
+        #advanced-btn {
+            font-size: .7rem !important;
+            line-height: 19px;
+            margin-top: 12px;
+            margin-bottom: 12px;
+            padding: 2px 8px;
+            border-radius: 14px !important;
+        }
+        #advanced-options {
+            display: none;
+            margin-bottom: 20px;
+        }
+        .footer {
+            margin-bottom: 45px;
+            margin-top: 35px;
+            text-align: center;
+            border-bottom: 1px solid #e5e5e5;
+        }
+        .footer>p {
+            font-size: .8rem;
+            display: inline-block;
+            padding: 0 10px;
+            transform: translateY(10px);
+            background: white;
+        }
+        .dark .footer {
+            border-color: #303030;
+        }
+        .dark .footer>p {
+            background: #0b0f19;
+        }
+        .acknowledgments h4{
+            margin: 1.25em 0 .25em 0;
+            font-weight: bold;
+            font-size: 115%;
+        }
+        #container-advanced-btns{
+            display: flex;
+            flex-wrap: wrap;
+            justify-content: space-between;
+            align-items: center;
+        }
+        .animate-spin {
+            animation: spin 1s linear infinite;
+        }
+        @keyframes spin {
+            from {
+                transform: rotate(0deg);
+            }
+            to {
+                transform: rotate(360deg);
+            }
+        }
+        #share-btn-container {
+            display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
+        }
+        #share-btn {
+            all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;
+        }
+        #share-btn * {
+            all: unset;
+        }
+        .gr-form{
+            flex: 1 1 50%; border-top-right-radius: 0; border-bottom-right-radius: 0;
+        }
+        #prompt-container{
+            gap: 0;
+        }
+        #generated_id{
+            min-height: 700px
+        }
+"""
+block = gr.Blocks(css=css)
+
+examples = [
+
+    [
+        'Thinking man in anime style'
+    ],
+
+]
+
+SPACE_ID = os.getenv('SPACE_ID')
+
+with block as demo:
+    gr.Markdown(f"""
+
+
+[![Framework: PyTorch](https://img.shields.io/badge/Framework-PyTorch-orange.svg)](https://pytorch.org/) [![Huggingface space](https://img.shields.io/badge/🤗-Huggingface-yello.svg)](https://huggingface.co/sberbank-ai/Kandinsky_2.0)
+
+<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>
+
+[Offical BlogPost](https://habr.com/ru/company/sberbank/blog/725282/)
+[Offical Telegram Bot](https://t.me/kandinsky21_bot)
+[Offical site](https://fusionbrain.ai/diffusion)
+
+## Model architecture:
+Kandinsky 2.1 inherits best practicies from Dall-E 2 and Latent diffusion, while introducing some new ideas.
+
+As text and image encoder it uses CLIP model and diffusion image prior (mapping) between latent spaces of CLIP modalities. This approach increases the visual performance of the model and unveils new horizons in blending images and text-guided image manipulation.
+
+For diffusion mapping of latent spaces we use transformer with num_layers=20, num_heads=32 and hidden_size=2048.
+
+Other architecture parts:
+
+- Text encoder (XLM-Roberta-Large-Vit-L-14) - 560M
+- Diffusion Image Prior — 1B
+- CLIP image encoder (ViT-L/14) - 427M
+- Latent Diffusion U-Net - 1.22B
+- MoVQ encoder/decoder - 67M
+
+Kandinsky 2.1 was trained on a large-scale image-text dataset LAION HighRes and fine-tuned on our internal datasets.
+
+**Kandinsky 2.1** architecture overview:
+
+![](https://raw.githubusercontent.com/ai-forever/Kandinsky-2/main/content/einstein.png)
+
+
+        """
+    )
+    with gr.Group():
+        with gr.Box():
+            with gr.Row().style(mobile_collapse=False, equal_height=True):
+
+                text = gr.Textbox(
+                    label="Enter your prompt", show_label=True, max_lines=2
+                ).style(
+                    border=(True, False, True, True),
+                    rounded=(True, False, False, True),
+                    container=False,
+                )
+                negative = gr.Textbox(
+                    label="Enter your negative prompt", show_label=True, max_lines=2
+                ).style(
+                    border=(True, False, True, True),
+                    rounded=(True, False, False, True),
+                    container=False,
+                )
+                btn = gr.Button("Run").style(
+                    margin=False,
+                    rounded=(False, True, True, False),
+                )
+
+        gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="generated_id").style(
+            grid=[2], height="auto"
+        )
+
+        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, negative], outputs=gallery, cache_examples=True)
+        ex.dataset.headers = [""]
+
+        text.submit(infer, inputs=[text, negative], outputs=gallery)
+        btn.click(infer, inputs=[text, negative], outputs=gallery)
+gr.Markdown("""
+
+
+# Authors
+
++ Arseniy Shakhmatov: [Github](https://github.com/cene555), [Blog](https://t.me/gradientdip)
++ Anton Razzhigaev: [Github](https://github.com/razzant), [Blog](https://t.me/abstractDL)
++ Aleksandr Nikolich: [Github](https://github.com/AlexWortega), [Blog](https://t.me/lovedeathtransformers)
++ Vladimir Arkhipkin: [Github](https://github.com/oriBetelgeuse)
++ Igor Pavlov: [Github](https://github.com/boomb0om)
++ Andrey Kuznetsov: [Github](https://github.com/kuznetsoffandrey)
++ Denis Dimitrov: [Github](https://github.com/denndimitrov)
+
+    """
+    )
+
+demo.queue(max_size=15).launch()