Add new model

model/model_manager.py

@@ -5,6 +5,7 @@ import requests
 import io, base64, json
 import spaces
 from PIL import Image
+from openai import OpenAI
 from .models import IMAGE_GENERATION_MODELS, IMAGE_EDITION_MODELS, load_pipeline
 from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_museum
 from serve.upload import get_random_mscoco_prompt
@@ -26,7 +27,7 @@ class ModelManager:
     @spaces.GPU(duration=120)
     def generate_image_ig(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
-        if 'cascade' not in name:
+        if 'Stable-cascade' not in name:
             result = pipe(prompt=prompt).images[0]
         else:
             prior, decoder = pipe
@@ -40,7 +41,6 @@ class ModelManager:
                 num_images_per_prompt=1,
                 num_inference_steps=20
             )
-
             decoder.enable_model_cpu_offload()
             result = decoder(
                 image_embeddings=prior_output.image_embeddings.to(torch.float16),
@@ -55,6 +55,7 @@ class ModelManager:
     def generate_image_ig_api(self, prompt, model_name):
         pipe = self.load_model_pipe(model_name)
         result = pipe(prompt=prompt)
+        
         return result
 
     def generate_image_ig_museum(self, model_name):

model/models/__init__.py

@@ -5,21 +5,8 @@ from .fal_api_models import load_fal_model
 from .huggingface_models import load_huggingface_model
 from .replicate_api_models import load_replicate_model
 from .openai_api_models import load_openai_model
+from .other_api_models import load_other_model
 
-# IMAGE_GENERATION_MODELS = ['huggingface_SD-v1.5_text2image',
-#                            'huggingface_SD-v2.1_text2image',
-#                            'huggingface_SD-XL-v1.0_text2image',
-#                            'huggingface_IF-I-XL-v1.0_text2image',
-#                             ]
-
-# IMAGE_GENERATION_MODELS = [ 'imagenhub_SD_generation',
-#                             'imagenhub_SDXL_generation',
-#                             'imagenhub_OpenJourney_generation',
-#                             'imagenhub_LCM_generation', 
-#                             'imagenhub_DeepFloydIF_generation', 
-#                             'imagenhub_PixArtAlpha_generation', 
-#                             'imagenhub_Kandinsky_generation', 
-#                             ]
 
 IMAGE_GENERATION_MODELS = [ 
                             'replicate_SDXL_text2image',
@@ -44,6 +31,11 @@ IMAGE_GENERATION_MODELS = [
                             'replicate_Deepfloyd-IF_text2image',
                             'huggingface_SD-turbo_text2image',
                             'huggingface_SDXL-turbo_text2image',
+                            'huggingface_Stable-cascade_text2image',
+                            'openai_Dalle-2_text2image',
+                            'openai_Dalle-3_text2image',
+                            'other_Midjourney-v6.0_text2image',
+                            'other_Midjourney-v5.0_text2image',
                             ]
 
 
@@ -78,7 +70,9 @@ def load_pipeline(model_name):
     elif model_source == "huggingface":
         pipe = load_huggingface_model(model_name, model_type)
     elif model_source == "openai":
-        pipe = load_openai_model(model_name)
+        pipe = load_openai_model(model_name, model_type)
+    elif model_source == "other":
+        pipe = load_other_model(model_name, model_type)
     else:
         raise ValueError(f"Model source {model_source} not supported")
     return pipe

model/models/huggingface_models.py

@@ -4,8 +4,6 @@ from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 import torch
 
 
-
-
 def load_huggingface_model(model_name, model_type):
     if model_name == "SD-turbo":
         pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float16, variant="fp16")
@@ -30,10 +28,10 @@ def load_huggingface_model(model_name, model_type):
 
 
 if __name__ == "__main__":
-    for name in ["SD-turbo", "SDXL-turbo"]: #"SD-turbo", "SDXL-turbo"
-        pipe = load_huggingface_model(name, "text2image")
-
+    # for name in ["SD-turbo", "SDXL-turbo"]: #"SD-turbo", "SDXL-turbo"
+    #     pipe = load_huggingface_model(name, "text2image")
 
     # for name in ["IF-I-XL-v1.0"]:
     #     pipe = load_huggingface_model(name, 'text2image')
     # pipe = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16)
+

model/models/openai_api_models.py

@@ -1,33 +1,59 @@
 from openai import OpenAI
+from PIL import Image
+import requests
+import io
+import os
+import base64
 
 
-def load_openai_model(model_name):
-    client = OpenAI()
+class OpenaiModel():
+    def __init__(self, model_name, model_type):
+        self.model_name = model_name
+        self.model_type = model_type
+    
+    def __call__(self, *args, **kwargs):
 
-    if model_name == "Dalle-3":
-        response = client.images.generate(
-            model="dall-e-3",
-            prompt="a white siamese cat",
-            size="1024x1024",
-            quality="standard",
-            n=1,
-            )
-    elif model_name == "Dalle-2":
-        response = client.images.generate(
-            model="dall-e-2",
-            prompt="a white siamese cat",
-            size="512x512",
-            quality="standard",
-            n=1,
-            )
-    else:
-        raise NotImplementedError
+        if self.model_type == "text2image":
+            assert "prompt" in kwargs, "prompt is required for text2image model"
 
-    image_url = response.data[0].url
+            client = OpenAI()
 
-    return image_url
+            if 'Dalle-3' in self.model_name:
+                client = OpenAI()
+                response = client.images.generate(
+                    model="dall-e-3",
+                    prompt=kwargs["prompt"],
+                    size="1024x1024",
+                    quality="standard",
+                    n=1,
+                    )
+            elif 'Dalle-2' in self.model_name:
+                client = OpenAI()
+                response = client.images.generate(
+                    model="dall-e-2",
+                    prompt=kwargs["prompt"],
+                    size="512x512",
+                    quality="standard",
+                    n=1,
+                    )
+            else:
+                raise NotImplementedError
+
+            result_url = response.data[0].url
+            response = requests.get(result_url)
+            result = Image.open(io.BytesIO(response.content))
+            return result
+        else:
+            raise ValueError("model_type must be text2image or image2image")
+
+
+
+def load_openai_model(model_name, model_type):
+    return OpenaiModel(model_name, model_type)
 
 
 if __name__ == "__main__":
-    image_url = load_openai_model('Dalle-3')
-    print(image_url)
+    pipe = load_openai_model('Dalle-2', 'text2image')
+    result = pipe(prompt='draw a tiger')
+    print(result)
+

model/models/other_api_models.py

@@ -0,0 +1,84 @@
+import requests
+import json
+import os
+from PIL import Image
+import io, time
+
+class OtherModel():
+    def __init__(self, model_name, model_type):
+        self.model_name = model_name
+        self.model_type = model_type
+        self.url = "https://www.xdai.online/mj/submit/imagine"
+        self.key = os.environ.get('MIDJOURNEY_KEY')
+        self.get_url = "https://www.xdai.online/mj/image/"
+        self.repeat_num = 5
+    
+    def __call__(self, *args, **kwargs):
+        if self.model_type == "text2image":
+            assert "prompt" in kwargs, "prompt is required for text2image model"
+            if self.model_name == "Midjourney-v6.0":
+                data = {
+                    "base64Array": [],
+                    "notifyHook": "",
+                    "prompt": "{} --v 6.0".format(kwargs["prompt"]),
+                    "state": "",
+                    "botType": "MID_JOURNEY",
+                }
+            elif self.model_name == "Midjourney-v5.0":
+                data = {
+                    "base64Array": [],
+                    "notifyHook": "",
+                    "prompt": "{} --v 5.0".format(kwargs["prompt"]),
+                    "state": "",
+                    "botType": "MID_JOURNEY",
+                }
+            else:
+                raise NotImplementedError
+            
+            headers = {
+                "Authorization": "Bearer {}".format(self.key),
+                "Content-Type": "application/json"
+            }
+            while 1:
+                response = requests.post(self.url, data=json.dumps(data), headers=headers)
+                if response.status_code == 200:
+                    print("Submit success!")
+                    response_json = json.loads(response.content.decode('utf-8'))
+                    img_id = response_json["result"]
+                    result_url = self.get_url + img_id
+                    self.repeat_num = 120
+                    while 1:
+                        time.sleep(1)
+                        img_response = requests.get(result_url)
+                        if img_response.status_code == 200:
+                            result = Image.open(io.BytesIO(img_response.content))
+                            width, height = result.size  
+                            new_width = width // 2  
+                            new_height = height // 2  
+                            result = result.crop((0, 0, new_width, new_height))  
+                            self.repeat_num = 5
+                            return result
+                        else:
+                            self.repeat_num = self.repeat_num - 1
+                            if self.repeat_num == 0:
+                                raise ValueError("Image request failed.")
+                            continue
+                        
+                else:
+                    self.repeat_num = self.repeat_num - 1
+                    if self.repeat_num == 0:
+                        raise ValueError("API request failed.")
+                    continue
+        else:
+            raise ValueError("model_type must be text2image")
+def load_other_model(model_name, model_type):
+    return OtherModel(model_name, model_type)
+
+if __name__ == "__main__":
+
+    pipe = load_other_model("Midjourney-v5.0", "text2image")
+    result = pipe(prompt="a good girl")
+    print(result)
+
+
+