Update handler.py

handler.py

@@ -26,101 +26,100 @@ num_return_sequences = 4  # the number of results to generate
 auto_mode = False
 
 prompt_modifier = PromptModifier(num_of_sequences=num_return_sequences)
-controlnet = ControlNet()
+
 lora_style = LoraStyle()
-text2img_pipe = Text2Img()
-img2img_pipe = Img2Img()
+
 slack = Slack()
 
 
 
-def get_patched_prompt(task: Task):
-    def add_style_and_character(prompt: List[str]):
-        for i in range(len(prompt)):
-            prompt[i] = add_code_names(prompt[i])
-            prompt[i] = lora_style.prepend_style_to_prompt(prompt[i], task.get_style())
+# def get_patched_prompt(task: Task):
+#     def add_style_and_character(prompt: List[str]):
+#         for i in range(len(prompt)):
+#             prompt[i] = add_code_names(prompt[i])
+#             prompt[i] = lora_style.prepend_style_to_prompt(prompt[i], task.get_style())
 
-    prompt = task.get_prompt()
+#     prompt = task.get_prompt()
 
-    if task.is_prompt_engineering():
-        prompt = prompt_modifier.modify(prompt)
-    else:
-        prompt = [prompt] * num_return_sequences
+#     if task.is_prompt_engineering():
+#         prompt = prompt_modifier.modify(prompt)
+#     else:
+#         prompt = [prompt] * num_return_sequences
 
-    ori_prompt = [task.get_prompt()] * num_return_sequences
+#     ori_prompt = [task.get_prompt()] * num_return_sequences
 
-    add_style_and_character(ori_prompt)
-    add_style_and_character(prompt)
+#     add_style_and_character(ori_prompt)
+#     add_style_and_character(prompt)
 
-    print({"prompts": prompt})
+#     print({"prompts": prompt})
 
-    return (prompt, ori_prompt)
+#     return (prompt, ori_prompt)
 
 
-# @update_db
-@auto_clear_cuda_and_gc(controlnet)
-@slack.auto_send_alert
-def canny(task: Task):
-    prompt, _ = get_patched_prompt(task)
+# # @update_db
+# @auto_clear_cuda_and_gc(controlnet)
+# @slack.auto_send_alert
+# def canny(task: Task):
+#     prompt, _ = get_patched_prompt(task)
 
-    controlnet.load_canny()
+#     controlnet.load_canny()
 
-    lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
-    lora_patcher.patch()
+#     lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
+#     lora_patcher.patch()
 
-    images = controlnet.process_canny(
-        prompt=prompt,
-        imageUrl=task.get_imageUrl(),
-        seed=task.get_seed(),
-        steps=task.get_steps(),
-        width=task.get_width(),
-        height=task.get_height(),
-        negative_prompt=[
-            f"monochrome, neon, x-ray, negative image, oversaturated, {task.get_negative_prompt()}"
-        ]
-        * num_return_sequences,
-        **lora_patcher.kwargs(),
-    )
+#     images = controlnet.process_canny(
+#         prompt=prompt,
+#         imageUrl=task.get_imageUrl(),
+#         seed=task.get_seed(),
+#         steps=task.get_steps(),
+#         width=task.get_width(),
+#         height=task.get_height(),
+#         negative_prompt=[
+#             f"monochrome, neon, x-ray, negative image, oversaturated, {task.get_negative_prompt()}"
+#         ]
+#         * num_return_sequences,
+#         **lora_patcher.kwargs(),
+#     )
 
-    generated_image_urls = upload_images(images, "_canny", task.get_taskId())
+#     generated_image_urls = upload_images(images, "_canny", task.get_taskId())
 
-    lora_patcher.cleanup()
-    controlnet.cleanup()
+#     lora_patcher.cleanup()
+#     controlnet.cleanup()
 
-    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+#     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
 
-# @update_db
-@auto_clear_cuda_and_gc(controlnet)
-@slack.auto_send_alert
-def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
-    prompt, _ = get_patched_prompt(task)
+# # @update_db
+# @auto_clear_cuda_and_gc(controlnet)
+# @slack.auto_send_alert
+# def pose(task: Task, s3_outkey: str = "_pose", poses: Optional[list] = None):
+#     prompt, _ = get_patched_prompt(task)
 
-    controlnet.load_pose()
+#     controlnet.load_pose()
 
-    lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
-    lora_patcher.patch()
+#     lora_patcher = lora_style.get_patcher(controlnet.pipe, task.get_style())
+#     lora_patcher.patch()
 
-    if poses is None:
-        poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
+#     if poses is None:
+#         poses = [controlnet.detect_pose(task.get_imageUrl())] * num_return_sequences
 
-    images = controlnet.process_pose(
-        prompt=prompt,
-        image=poses,
-        seed=task.get_seed(),
-        steps=task.get_steps(),
-        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
-        width=task.get_width(),
-        height=task.get_height(),
-        **lora_patcher.kwargs(),
-    )
+#     images = controlnet.process_pose(
+#         prompt=prompt,
+#         image=poses,
+#         seed=task.get_seed(),
+#         steps=task.get_steps(),
+#         negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+#         width=task.get_width(),
+#         height=task.get_height(),
+#         **lora_patcher.kwargs(),
+#     )
 
-    generated_image_urls = upload_images(images, s3_outkey, task.get_taskId())
+#     generated_image_urls = upload_images(images, s3_outkey, task.get_taskId())
 
-    lora_patcher.cleanup()
-    controlnet.cleanup()
+#     lora_patcher.cleanup()
+#     controlnet.cleanup()
 
-    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+#     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
 
 # @update_db
@@ -153,30 +152,30 @@ def text2img(task: Task, text2img_pipe ):
     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
 
-# @update_db
-@auto_clear_cuda_and_gc(controlnet)
-@slack.auto_send_alert
-def img2img(task: Task):
-    prompt, _ = get_patched_prompt(task)
+# # @update_db
+# @auto_clear_cuda_and_gc(controlnet)
+# @slack.auto_send_alert
+# def img2img(task: Task):
+#     prompt, _ = get_patched_prompt(task)
 
-    lora_patcher = lora_style.get_patcher(img2img_pipe.pipe, task.get_style())
-    lora_patcher.patch()
+#     lora_patcher = lora_style.get_patcher(img2img_pipe.pipe, task.get_style())
+#     lora_patcher.patch()
 
-    torch.manual_seed(task.get_seed())
+#     torch.manual_seed(task.get_seed())
 
-    images = img2img_pipe.process(
-        prompt=prompt,
-        imageUrl=task.get_imageUrl(),
-        negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
-        steps=task.get_steps(),
-        **lora_patcher.kwargs(),
-    )
+#     images = img2img_pipe.process(
+#         prompt=prompt,
+#         imageUrl=task.get_imageUrl(),
+#         negative_prompt=[task.get_negative_prompt()] * num_return_sequences,
+#         steps=task.get_steps(),
+#         **lora_patcher.kwargs(),
+#     )
 
-    generated_image_urls = upload_images(images, "_imgtoimg", task.get_taskId())
+#     generated_image_urls = upload_images(images, "_imgtoimg", task.get_taskId())
 
-    lora_patcher.cleanup()
+#     lora_patcher.cleanup()
 
-    return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
+#     return {"modified_prompts": prompt, "generated_image_urls": generated_image_urls}
 
 
 
@@ -213,14 +212,23 @@ class EndpointHandler():
             # self.multi_controlnet_model[model["model_id"]] = controlnet.load(model["model_id"])
             # self.multi_text2image_model[model["model_id"]] = text2img_pipe.load(model["model_id"])
             # self.multi_image2image_model[model["model_id"]] = img2img_pipe.load(model["model_id"])
-            self.multi_controlnet_model[model["model_id"]] = controlnet.load(model["model_id"])
-            self.multi_text2image_model[model["model_id"]] = text2img_pipe.load( model["model_id"])
-            self.multi_image2image_model[model["model_id"]] = img2img_pipe.load( model["model_id"])
+            controlnet = ControlNet()
+            img2img_pipe = Img2Img()
+            text2img_pipe = Text2Img()
+            self.multi_controlnet_model[model["model_id"]] = controlnet;
+            controlnet.load(model["model_id"])
+            
+            
+            self.multi_text2image_model[model["model_id"]] = text2img_pipe;
+            text2img_pipe.load( model["model_id"])
+            
+            self.multi_image2image_model[model["model_id"]] = img2img_pipe;
+            img2img_pipe.load( model["model_id"])
             
             print(" Logs: model[model_id]", model["model_id"])
-            print("Logs: multimodel controlnet pipelines are",  path + model["model_id"])    
-            print("Logs: multimodel text2img pipelines are",  path + model["model_id"]) 
-            print("Logs: multimodel imgtoimage pipelines are",  path + model["model_id"]) 
+            print("Logs: multimodel controlnet pipelines are",  self.multi_controlnet_model[model["model_id"]])    
+            print("Logs: multimodel text2img pipelines are",  self.multi_text2image_model[model["model_id"]]) 
+            print("Logs: multimodel imgtoimage pipelines are",  self.multi_image2image_model[model["model_id"]]) 
         # controlnet.load(path)
         # text2img_pipe.load(path)
         # img2img_pipe.load(path)
@@ -274,15 +282,16 @@ class EndpointHandler():
             if task_type == TaskType.TEXT_TO_IMAGE:
                 #  character sheet
                 if "character sheet" in task.get_prompt().lower():
-                    return pose(task, s3_outkey="", poses=pickPoses())
+                    print("pose is here")
+                    # return pose(task, s3_outkey="", poses=pickPoses())
                 else:
                     return text2img(task, self.multi_text2image_model[model_id])
-            elif task_type == TaskType.IMAGE_TO_IMAGE:
-                return img2img(task)
-            elif task_type == TaskType.CANNY:
-                return canny(task)
-            elif task_type == TaskType.POSE:
-                return pose(task)
+            # elif task_type == TaskType.IMAGE_TO_IMAGE:
+            #     return img2img(task)
+            # elif task_type == TaskType.CANNY:
+            #     return canny(task)
+            # elif task_type == TaskType.POSE:
+            #     return pose(task)
             else:
                 raise Exception("Invalid task type")
         except Exception as e: