Add pose detection, controlnet using pose

app.py

@@ -25,6 +25,7 @@ import cv2
 import einops
 from pytorch_lightning import seed_everything
 import random
+from controlnet_aux import OpenposeDetector
 
 VISUAL_CHATGPT_PREFIX = """Visual ChatGPT is designed to be able to assist with a wide range of text and visual related tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. Visual ChatGPT is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.
 
@@ -272,6 +273,61 @@ class canny2image:
         real_image.save(updated_image_path)
         return updated_image_path
 
+class image2pose:
+    def __init__(self):
+        print("Direct detect pose.")
+        
+    def inference(self, inputs):
+        print("===>Starting image2pose Inference")
+        image = Image.open(inputs)
+        image = np.array(image)
+        pose_image = pose_model(image)
+        
+        updated_image_path = get_new_image_name(inputs, func_name="pose")
+        pose_image.save(updated_image_path)
+        return updated_image_path
+
+
+class pose2image:
+    def __init__(self, device):
+        print("Initialize the pose2image model.")
+        
+
+        # Models
+        controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16)
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
+        )
+        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+        
+        # This command loads the individual model components on GPU on-demand. So, we don't
+        # need to explicitly call pipe.to("cuda").
+        self.pipe.enable_model_cpu_offload()
+        
+        self.pipe.enable_xformers_memory_efficient_attention()
+        
+        # Generator seed,
+        self.generator = torch.manual_seed(0)
+    
+
+    def get_pose(self,image):
+        return pose_model(image)
+        
+    def inference(self, inputs):
+        print("===>Starting pose2image Inference")
+        image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
+        image = Image.open(image_path)
+        image = np.array(image)
+        prompt = instruct_text
+        pose_image = self.get_pose(image)
+        output = self.pipe(prompt,pose_image,generator=self.generator,num_images_per_prompt=1,num_inference_steps=20)
+        
+        updated_image_path = get_new_image_name(image_path, func_name="pose2image")
+        real_image = output.images[0]  # get default the index0 image
+        real_image.save(updated_image_path)
+        return updated_image_path
+
+
 class BLIPVQA:
     def __init__(self, device):
         print("Initializing BLIP VQA to %s" % device)
@@ -296,6 +352,8 @@ class ConversationBot:
         self.t2i = T2I(device="cuda:0")
         self.image2canny = image2canny()
         self.canny2image = canny2image(device="cuda:0")
+        self.image2pose = image2pose()
+        self.pose2image = pose2image(device="cuda:0")
         self.BLIPVQA = BLIPVQA(device="cuda:0")
         #self.pix2pix = Pix2Pix(device="cuda:0")
         self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
@@ -361,12 +419,12 @@ class ConversationBot:
             #Tool(name="Generate Image Condition On Sketch Image", func=self.scribble2image.inference,
                  #description="useful when you want to generate a new real image from both the user desciption and a scribble image or a sketch image. "
                              #"The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
-            #Tool(name="Pose Detection On Image", func=self.image2pose.inference,
-                 #description="useful when you want to detect the human pose of the image. like: generate human poses of this image, or generate a pose image from this image. "
-                             #"The input to this tool should be a string, representing the image_path"),
-            #Tool(name="Generate Image Condition On Pose Image", func=self.pose2image.inference,
-                 #description="useful when you want to generate a new real image from both the user desciption and a human pose image. like: generate a real image of a human from this human pose image, or generate a new real image of a human from this pose. "
-                             #"The input to this tool should be a comma seperated string of two, representing the image_path and the user description")]
+            Tool(name="Pose Detection On Image", func=self.image2pose.inference,
+                 description="useful when you want to detect the human pose of the image. like: generate human poses of this image, or generate a pose image from this image. "
+                             "The input to this tool should be a string, representing the image_path"),
+            Tool(name="Generate Image Condition On Pose Image", func=self.pose2image.inference,
+                 description="useful when you want to generate a new real image from both the user desciption and a human pose image. like: generate a real image of a human from this human pose image, or generate a new real image of a human from this pose. "
+                             "The input to this tool should be a comma seperated string of two, representing the image_path and the user description")]
         ]
 
     def init_langchain(self,api_key):