update

model/model_manager.py

@@ -1,68 +1,28 @@
 import concurrent.futures 
 import random
 import gradio as gr
-# from fal_api_utils import load_fal_model
-from .imagenhub_utils import load_imagenhub_model
-import spaces
 import requests
 import io, base64, json
 from PIL import Image
-import os
-
-
-IMAGE_GENERATION_MODELS = ['imagenhub_LCM_generation','imagenhub_SDXLTurbo_generation','imagenhub_SDXL_generation', 'imagenhub_PixArtAlpha_generation',
-                            'imagenhub_OpenJourney_generation','imagenhub_SDXLLightning_generation', 'imagenhub_StableCascade_generation',
-                            'imagenhub_PlayGroundV2_generation', 'imagenhub_PlayGroundV2.5_generation']
-IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
-                        'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition']
+from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, load_pipeline
 
 class ModelManager:
     def __init__(self):
         self.model_ig_list = IMAGE_GENERATION_MODELS
         self.model_ie_list = IMAGE_EDITION_MODELS
         self.loaded_models = {}
-    # @spaces.GPU
+
     def load_model_pipe(self, model_name):
-        model_source, model_name, model_type = model_name.split("_")
         if not model_name in self.loaded_models:
-            if model_source == "imagenhub":
-                pipe = load_imagenhub_model(model_name)
-            # elif model_source == "fal":
-            #     pipe = load_fal_model(model_name, model_type)
-            else:
-                raise ValueError(f"Model source {model_source} not supported")
+            pipe = load_pipeline(model_name)
             self.loaded_models[model_name] = pipe
         else:
             pipe = self.loaded_models[model_name]
         return pipe
-    
-    def generate_image_playground(self, model_name, prompt):
-        if model_name == "imagenhub_PlayGroundV2_generation":
-            model_name = "Playground_v2"
-        elif model_name == "imagenhub_PlayGroundV2.5_generation":
-            model_name = "Playground_v2.5"
-
-        headers = {
-            'Content-Type': 'application/json',
-            'Authorization': os.environ['PlaygroundAPI'],
-        }
-
-        data = json.dumps({"prompt": prompt, "filter_model": model_name, "scheduler": "DPMPP_2M_K", "guidance_scale": 3})
 
-        response = requests.post('https://playground.com/api/models/external/v1', headers=headers, data=data)
-        response.raise_for_status() 
-        json_obj = response.json()
-        image_base64 = json_obj['images'][0]
-        img = Image.open(io.BytesIO(base64.decodebytes(bytes(image_base64, "utf-8"))))
-
-        return img
-    @spaces.GPU(duration=60)
     def generate_image_ig(self, prompt, model_name):
-        if 'playground' in model_name.lower():
-            result = self.generate_image_playground(model_name=model_name, prompt=prompt)
-        else:
-            pipe = self.load_model_pipe(model_name)
-            result = pipe(prompt=prompt)
+        pipe = self.load_model_pipe(model_name)
+        result = pipe(prompt=prompt)
         return result
 
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B):
@@ -88,7 +48,7 @@ class ModelManager:
                 result = future.result()
                 results.append(result)
         return results[0], results[1]
-    @spaces.GPU(duration=150)
+
     def generate_image_ie(self, textbox_source, textbox_target, textbox_instruct, source_image, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(src_image = source_image, src_prompt = textbox_source, target_prompt = textbox_target, instruct_prompt = textbox_instruct)

model/models/__init__.py

@@ -0,0 +1,27 @@
+from .imagenhub_models import load_imagenhub_model
+from .playground_api import load_playground_model
+
+IMAGE_GENERATION_MODELS = ['imagenhub_LCM_generation','imagenhub_SDXLTurbo_generation','imagenhub_SDXL_generation', 'imagenhub_PixArtAlpha_generation',
+                            'imagenhub_OpenJourney_generation','imagenhub_SDXLLightning_generation', 'imagenhub_StableCascade_generation',
+                            'playground_PlayGroundV2_generation', 'playground_PlayGroundV2.5_generation']
+IMAGE_EDITION_MODELS = ['imagenhub_CycleDiffusion_edition', 'imagenhub_Pix2PixZero_edition', 'imagenhub_Prompt2prompt_edition',
+                        'imagenhub_SDEdit_edition', 'imagenhub_InstructPix2Pix_edition', 'imagenhub_MagicBrush_edition', 'imagenhub_PNP_edition']
+
+
+def load_pipeline(model_name):
+    """
+    Load a model pipeline based on the model name
+    Args:
+        model_name (str): The name of the model to load, should be of the form {source}_{name}_{type}
+        the source can be either imagenhub or playground
+        the name is the name of the model used to load the model
+        the type is the type of the model, either generation or edition
+    """
+    model_source, model_name, model_type = model_name.split("_")
+    if model_source == "imagenhub":
+        pipe = load_imagenhub_model(model_name, model_type)
+    elif model_source == "playground":
+        pipe = load_playground_model(model_name)
+    else:
+        raise ValueError(f"Model source {model_source} not supported")
+    return pipe

model/models/playground_api.py

@@ -0,0 +1,35 @@
+import os
+import json
+import requests
+from PIL import Image
+import io
+import base64
+class PlayGround():
+    def __init__(self, model_name, model_type=None):
+        self.model_name = model_name
+        self.model_type = model_type
+        self.api_key = os.environ['PlaygroundAPI']
+        if model_name == "PlayGroundV2":
+            self._model_name = "Playground_v2"
+        elif model_name == "PlayGroundV2.5":
+            self._model_name = "Playground_v2.5"
+
+        
+    def __call__(self, prompt):
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': "Bearer " + self.api_key,
+        }
+
+        data = json.dumps({"prompt": prompt, "filter_model": self._model_name, "scheduler": "DPMPP_2M_K", "guidance_scale": 3})
+
+        response = requests.post('https://playground.com/api/models/external/v1', headers=headers, data=data)
+        response.raise_for_status() 
+        json_obj = response.json()
+        image_base64 = json_obj['images'][0]
+        img = Image.open(io.BytesIO(base64.decodebytes(bytes(image_base64, "utf-8"))))
+
+        return img
+    
+def load_playground_model(model_name, model_type="generation"):
+    return PlayGround(model_name, model_type)

serve/gradio_web.py

@@ -46,9 +46,9 @@ Find out who is the 🥇conditional image generation models! More models are goi
             gr.Markdown(model_description_md, elem_id="model_description_markdown")
         with gr.Row():
             with gr.Column():
-                chatbot_left = gr.Image(height=512, width=512, label = "Model A")
+                chatbot_left = gr.Image(width=512, label = "Model A")
             with gr.Column():
-                chatbot_right = gr.Image(height=512, width=512, label = "Model B")
+                chatbot_right = gr.Image(width=512, label = "Model B")
 
         with gr.Row():
             with gr.Column():
@@ -86,7 +86,7 @@ Find out who is the 🥇conditional image generation models! More models are goi
 
     gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
 
-    dummy_img_output = gr.Image(height=512, width=512, visible=False)
+    dummy_img_output = gr.Image(width=512, visible=False)
     gr.Examples(
         examples=[["a cute dog is playing a ball", os.path.join("./examples", "dog.jpg")],
                   ["Buildings on fire, old film still", os.path.join("./examples", "fire.jpg")],
@@ -215,9 +215,9 @@ def build_side_by_side_ui_named(models):
 
         with gr.Row():
             with gr.Column():
-                chatbot_left = gr.Image(height=512, width=512, label = "Model A")
+                chatbot_left = gr.Image(width=512, label = "Model A")
             with gr.Column():
-                chatbot_right = gr.Image(height=512, width=512, label = "Model B")
+                chatbot_right = gr.Image(width=512, label = "Model B")
         with gr.Row():
             leftvote_btn = gr.Button(
                 value="👈  A is better", visible=False, interactive=False
@@ -245,7 +245,7 @@ def build_side_by_side_ui_named(models):
 
     gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
 
-    dummy_img_output = gr.Image(height=512, width=512, visible=False)
+    dummy_img_output = gr.Image(width=512, visible=False)
     gr.Examples(
         examples=[["a cute dog is playing a ball", os.path.join("./examples", "dog.jpg")],
                   ["Buildings on fire, old film still", os.path.join("./examples", "fire.jpg")],
@@ -378,7 +378,7 @@ def build_single_model_ui(models, add_promotion_links=False):
         send_btn = gr.Button(value="Send", variant="primary", scale=0)
 
     with gr.Row():
-        chatbot = gr.Image(height=512, width=512)
+        chatbot = gr.Image()
 
     with gr.Row() as button_row:
         upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
@@ -390,7 +390,7 @@ def build_single_model_ui(models, add_promotion_links=False):
     if add_promotion_links:
         gr.Markdown(acknowledgment_md, elem_id="ack_markdown")
 
-    dummy_img_output = gr.Image(height=512, width=512, visible=False)
+    dummy_img_output = gr.Image(visible=False)
     gr.Examples(
         examples=[["a cute dog is playing a ball", os.path.join("./examples", "dog.jpg")],
                   ["Buildings on fire, old film still", os.path.join("./examples", "fire.jpg")],

serve/gradio_web_image_editing.py

@@ -48,9 +48,9 @@ Find out who is the 🥇conditional image edition models!
             gr.Markdown(model_description_md, elem_id="model_description_markdown")
         with gr.Row():
             with gr.Column():
-                    chatbot_left = gr.Image(height=512, width=512, type="pil")
+                    chatbot_left = gr.Image(width=512, type="pil")
             with gr.Column():
-                    chatbot_right = gr.Image(height=512, width=512, type="pil")
+                    chatbot_right = gr.Image(width=512, type="pil")
 
         with gr.Row():
             with gr.Column():
@@ -249,9 +249,9 @@ def build_side_by_side_ui_named_ie(models):
 
         with gr.Row():
             with gr.Column():
-                chatbot_left = gr.Image(height=512, width=512, label = "Model A")
+                chatbot_left = gr.Image(width=512, label = "Model A")
             with gr.Column():
-                chatbot_right = gr.Image(height=512, width=512, label = "Model B")
+                chatbot_right = gr.Image(width=512, label = "Model B")
 
         with gr.Row():
             leftvote_btn = gr.Button(

serve/vote_utils.py

@@ -292,9 +292,9 @@ class ImageStateIE:
 
 def generate_ig(gen_func, state, text, model_name, request: gr.Request):
     if not text:
-        raise gr.Error("Prompt cannot be empty.")
+        raise gr.Warning("Prompt cannot be empty.")
     if not model_name:
-        raise gr.Error("Model name cannot be empty.")
+        raise gr.Warning("Model name cannot be empty.")
     if state is None:
         state = ImageStateIG(model_name)
     ip = get_ip(request)
@@ -332,11 +332,11 @@ def generate_ig(gen_func, state, text, model_name, request: gr.Request):
         
 def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):
     if not text:
-        raise gr.Error("Prompt cannot be empty.")
+        raise gr.Warning("Prompt cannot be empty.")
     if not model_name0:
-        raise gr.Error("Model name A cannot be empty.")
+        raise gr.Warning("Model name A cannot be empty.")
     if not model_name1:
-        raise gr.Error("Model name B cannot be empty.")
+        raise gr.Warning("Model name B cannot be empty.")
     if state0 is None:
         state0 = ImageStateIG(model_name0)
     if state1 is None:
@@ -392,7 +392,7 @@ def generate_igm(gen_func, state0, state1, text, model_name0, model_name1, reque
             
 def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1, request: gr.Request):
     if not text:
-        raise gr.Error("Prompt cannot be empty.")
+        raise gr.Warning("Prompt cannot be empty.")
     if state0 is None:
         state0 = ImageStateIG(model_name0)
     if state1 is None:
@@ -450,15 +450,15 @@ def generate_igm_annoy(gen_func, state0, state1, text, model_name0, model_name1,
             
 def generate_ie(gen_func, state, source_text, target_text, instruct_text, source_image, model_name, request: gr.Request):
     if not source_text:
-        raise gr.Error("Source prompt cannot be empty.")
+        raise gr.Warning("Source prompt cannot be empty.")
     if not target_text:
-        raise gr.Error("Target prompt cannot be empty.")
+        raise gr.Warning("Target prompt cannot be empty.")
     if not instruct_text:
-        raise gr.Error("Instruction prompt cannot be empty.")
+        raise gr.Warning("Instruction prompt cannot be empty.")
     if not source_image:
-        raise gr.Error("Source image cannot be empty.")
+        raise gr.Warning("Source image cannot be empty.")
     if not model_name:
-        raise gr.Error("Model name cannot be empty.")
+        raise gr.Warning("Model name cannot be empty.")
     if state is None:
         state = ImageStateIE(model_name)
     ip = get_ip(request)
@@ -503,17 +503,17 @@ def generate_ie(gen_func, state, source_text, target_text, instruct_text, source
         
 def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):
     if not source_text:
-        raise gr.Error("Source prompt cannot be empty.")
+        raise gr.Warning("Source prompt cannot be empty.")
     if not target_text:
-        raise gr.Error("Target prompt cannot be empty.")
+        raise gr.Warning("Target prompt cannot be empty.")
     if not instruct_text:
-        raise gr.Error("Instruction prompt cannot be empty.")
+        raise gr.Warning("Instruction prompt cannot be empty.")
     if not source_image:
-        raise gr.Error("Source image cannot be empty.")
+        raise gr.Warning("Source image cannot be empty.")
     if not model_name0:
-        raise gr.Error("Model name A cannot be empty.")
+        raise gr.Warning("Model name A cannot be empty.")
     if not model_name1:
-        raise gr.Error("Model name B cannot be empty.")
+        raise gr.Warning("Model name B cannot be empty.")
     if state0 is None:
         state0 = ImageStateIE(model_name0)
     if state1 is None:
@@ -580,13 +580,13 @@ def generate_iem(gen_func, state0, state1, source_text, target_text, instruct_te
 
 def generate_iem_annoy(gen_func, state0, state1, source_text, target_text, instruct_text, source_image, model_name0, model_name1, request: gr.Request):
     if not source_text:
-        raise gr.Error("Source prompt cannot be empty.")
+        raise gr.Warning("Source prompt cannot be empty.")
     if not target_text:
-        raise gr.Error("Target prompt cannot be empty.")
+        raise gr.Warning("Target prompt cannot be empty.")
     if not instruct_text:
-        raise gr.Error("Instruction prompt cannot be empty.")
+        raise gr.Warning("Instruction prompt cannot be empty.")
     if not source_image:
-        raise gr.Error("Source image cannot be empty.")
+        raise gr.Warning("Source image cannot be empty.")
     if state0 is None:
         state0 = ImageStateIE(model_name0)
     if state1 is None: