Upload multit2i.py

multit2i.py

@@ -104,7 +104,7 @@ def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
 def load_model(model_name: str):
     global loaded_models
     global model_info_dict
-    if model_name in loaded_models.keys(): return model_name
+    if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
         loaded_models[model_name] = gr.load(f'models/{model_name}')
         print(f"Loaded: {model_name}")
@@ -112,13 +112,13 @@ def load_model(model_name: str):
         if model_name in loaded_models.keys(): del loaded_models[model_name]
         print(f"Failed to load: {model_name}")
         print(e)
-        return ""
+        return None
     try:
         model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
     except Exception as e:
         if model_name in model_info_dict.keys(): del model_info_dict[model_name]
         print(e)
-    return model_name
+    return loaded_models[model_name]
 
 
 async def async_load_models(models: list, limit: int=5):
@@ -163,12 +163,12 @@ def infer(prompt: str, model_name: str, recom_prompt: bool, progress=gr.Progress
     caption = model_name.split("/")[-1]
     try:
         model = load_model(model_name)
-        if not model: return ("", None)
+        if not model: return (Image.Image(), None)
         image_path = model(prompt + rprompt + seed)
         image = Image.open(image_path).convert('RGB')
     except Exception as e:
         print(e)
-        return ("", None)
+        return (Image.Image(), None)
     return (image, caption)