Update model/model_manager.py (#6)

- Update model/model_manager.py (5c6a38a87ac0282c7bee84f4a76ac55afe06e266)


Co-authored-by: Yuansheng Ni <[email protected]>

model/model_manager.py

@@ -21,7 +21,7 @@ class ModelManager:
         self.model_ig_list = IMAGE_GENERATION_MODELS
         self.model_ie_list = IMAGE_EDITION_MODELS
         self.loaded_models = {}
-    @spaces.GPU
+    # @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:
@@ -56,7 +56,7 @@ class ModelManager:
         img = Image.open(io.BytesIO(base64.decodebytes(bytes(image_base64, "utf-8"))))
 
         return img
-    @spaces.GPU
+    @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)
@@ -64,7 +64,7 @@ class ModelManager:
             pipe = self.load_model_pipe(model_name)
             result = pipe(prompt=prompt)
         return result
-    # @spaces.GPU
+
     def generate_image_ig_parallel_anony(self, prompt, model_A, model_B):
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in self.model_ig_list], 2)
@@ -78,7 +78,7 @@ class ModelManager:
                 result = future.result()
                 results.append(result)
         return results[0], results[1], model_names[0], model_names[1]
-    # @spaces.GPU
+
     def generate_image_ig_parallel(self, prompt, model_A, model_B):
         results = []
         model_names = [model_A, model_B]
@@ -88,12 +88,12 @@ class ModelManager:
                 result = future.result()
                 results.append(result)
         return results[0], results[1]
-    @spaces.GPU
+    @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)
         return result
-    # @spaces.GPU
+
     def generate_image_ie_parallel(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         results = []
         model_names = [model_A, model_B]
@@ -103,7 +103,7 @@ class ModelManager:
                 result = future.result()
                 results.append(result)
         return results[0], results[1]
-    # @spaces.GPU
+
     def generate_image_ie_parallel_anony(self, textbox_source, textbox_target, textbox_instruct, source_image, model_A, model_B):
         if model_A == "" and model_B == "":
             model_names = random.sample([model for model in self.model_ie_list], 2)