Update app.py

app.py

@@ -2,8 +2,10 @@ import gradio as gr
 import torch
 from torch import autocast
 from kandinsky2 import get_kandinsky2
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = get_kandinsky2('cuda', 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, "Ultra": 150}
     length = length_dict[quality]
@@ -19,4 +21,9 @@ iface = gr.Interface(
   outputs=gr.outputs.Image(label="Generated image:")
 )
 
-iface.launch()
+if device.type == 'cpu':
+    model.load_state_dict(torch.load('path/to/model.pth', map_location=device))
+else:
+    model.load_state_dict(torch.load('path/to/model.pth'))
+
+iface.launch()