Update app.py

app.py

@@ -1,11 +1,14 @@
 import gradio as gr
-from transformers import pipeline
-text_generator = pipeline(model="ai-forever/Kandinsky_2.1", task="text2img")
-
+from kandinsky2 import get_kandinsky2
+model = get_kandinsky2('cpu', task_type='text2img', cache_dir='/tmp/kandinsky2', model_version='2.1', use_flash_attention=False)
 def generate_text(prompt, quality="High (Default)"):
-    length_dict = {"Low": 50, "High (Default)": 100, "High": 175}
+    length_dict = {"Low": 50, "High (Default)": 100, "Ultra": 150}
     length = length_dict[quality]
-    return text_generator(prompt, h=768, w=768, num_steps=legnth, batch_size=1, sampler='p_sampler', prior_steps="5", prior_cf_scale=4,)[0]
+    return model.generate_text2img('''red cat, 4k photo''', num_steps=length,
+                          batch_size=1, guidance_scale=4,
+                           h=768, w=768
+                           ,sampler='p_sampler', prior_cf_scale=4,
+                            prior_steps="5",)[0]
 
 iface = gr.Interface(
   fn=generate_text,