envs

app.py

@@ -295,7 +295,7 @@ class ImageConductor:
         if isinstance(tracking_points, list):
             input_all_points = tracking_points
         else:
-            input_all_points = tracking_points.constructor_args['value']
+            input_all_points = tracking_points
         
         
         resized_all_points = [tuple([tuple([float(e1[0]*self.width/original_width), float(e1[1]*self.height/original_height)]) for e1 in e]) for e in input_all_points]
@@ -304,6 +304,10 @@ class ImageConductor:
         id = base.split('_')[-1]
         
         
+        # with open(f'{output_dir}/points-{id}.json', 'w') as f:
+        #     json.dump(input_all_points, f)
+        
+        
         visualized_drag, _ = visualize_drag(first_frame_path, resized_all_points, drag_mode, self.width, self.height, self.model_length)
 
         ## image condition        
@@ -377,16 +381,18 @@ class ImageConductor:
         vis_video = (rearrange(sample[0], 'c t h w -> t h w c') * 255.).clip(0, 255)
         torchvision.io.write_video(outputs_path, vis_video, fps=8, video_codec='h264', options={'crf': '10'})
 
-        return visualized_drag, outputs_path
+        return {output_image: visualized_drag, output_video: outputs_path}
 
 
 def reset_states(first_frame_path, tracking_points):
     first_frame_path = gr.State()
     tracking_points = gr.State([])
-    return None, first_frame_path, tracking_points
+    return {input_image:None, first_frame_path_var: first_frame_path, tracking_points_var: tracking_points}
 
 
-def preprocess_image(image):
+def preprocess_image(image, tracking_points):
+    if len(tracking_points) != 0:
+        tracking_points = gr.State([])
     image_pil = image2pil(image.name)
     raw_w, raw_h = image_pil.size
     resize_ratio = max(384/raw_w, 256/raw_h)
@@ -395,7 +401,7 @@ def preprocess_image(image):
     id = str(uuid.uuid4())[:4]
     first_frame_path = os.path.join(output_dir, f"first_frame_{id}.jpg")
     image_pil.save(first_frame_path, quality=95)
-    return first_frame_path, first_frame_path, gr.State([])
+    return {input_image: first_frame_path, first_frame_path_var: first_frame_path, tracking_points_var: tracking_points}
 
 
 def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.SelectData):  # SelectData is a subclass of EventData
@@ -405,13 +411,13 @@ def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.Se
         color = (0, 0, 255, 255)
 
     print(f"You selected {evt.value} at {evt.index} from {evt.target}")
-    tracking_points.constructor_args['value'][-1].append(evt.index)
-    print(tracking_points.constructor_args)
+    tracking_points[-1].append(evt.index)
+    print(tracking_points)
     
     transparent_background = Image.open(first_frame_path).convert('RGBA')
     w, h = transparent_background.size
     transparent_layer = np.zeros((h, w, 4))
-    for track in tracking_points.constructor_args['value']:
+    for track in tracking_points:
         if len(track) > 1:
             for i in range(len(track)-1):
                 start_point = track[i]
@@ -428,13 +434,13 @@ def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.Se
 
     transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
     trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-    return tracking_points, trajectory_map
+    return {tracking_points_var: tracking_points, input_image: trajectory_map}
 
 
 def add_drag(tracking_points):
-    tracking_points.constructor_args['value'].append([])
-    print(tracking_points.constructor_args)
-    return tracking_points
+    tracking_points.append([])
+    print(tracking_points)
+    return {tracking_points_var: tracking_points}
     
 
 def delete_last_drag(tracking_points, first_frame_path, drag_mode):
@@ -442,11 +448,11 @@ def delete_last_drag(tracking_points, first_frame_path, drag_mode):
         color = (255, 0, 0, 255)
     elif drag_mode=='camera':
         color = (0, 0, 255, 255)
-    tracking_points.constructor_args['value'].pop()
+    tracking_points.pop()
     transparent_background = Image.open(first_frame_path).convert('RGBA')
     w, h = transparent_background.size
     transparent_layer = np.zeros((h, w, 4))
-    for track in tracking_points.constructor_args['value']:
+    for track in tracking_points:
         if len(track) > 1:
             for i in range(len(track)-1):
                 start_point = track[i]
@@ -463,7 +469,7 @@ def delete_last_drag(tracking_points, first_frame_path, drag_mode):
 
     transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
     trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-    return tracking_points, trajectory_map
+    return {tracking_points_var: tracking_points, input_image: trajectory_map}
     
 
 def delete_last_step(tracking_points, first_frame_path, drag_mode):
@@ -471,11 +477,11 @@ def delete_last_step(tracking_points, first_frame_path, drag_mode):
         color = (255, 0, 0, 255)
     elif drag_mode=='camera':
         color = (0, 0, 255, 255)
-    tracking_points.constructor_args['value'][-1].pop()
+    tracking_points[-1].pop()
     transparent_background = Image.open(first_frame_path).convert('RGBA')
     w, h = transparent_background.size
     transparent_layer = np.zeros((h, w, 4))
-    for track in tracking_points.constructor_args['value']:
+    for track in tracking_points:
         if len(track) > 1:
             for i in range(len(track)-1):
                 start_point = track[i]
@@ -492,147 +498,148 @@ def delete_last_step(tracking_points, first_frame_path, drag_mode):
 
     transparent_layer = Image.fromarray(transparent_layer.astype(np.uint8))
     trajectory_map = Image.alpha_composite(transparent_background, transparent_layer)
-    return tracking_points, trajectory_map
-
-
-block = gr.Blocks(
-        theme=gr.themes.Soft(
-             radius_size=gr.themes.sizes.radius_none,
-             text_size=gr.themes.sizes.text_md
-         )
-        ).queue()
-with block as demo:
-    with gr.Row():
-        with gr.Column():
-            gr.HTML(head)
-
-    gr.Markdown(descriptions)
-
-    with gr.Accordion(label="🛠️ Instructions:", open=True, elem_id="accordion"):
-        with gr.Row(equal_height=True):
-            gr.Markdown(instructions)      
-
-
-    # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    device = torch.device("cuda")
-    unet_path = 'models/unet.ckpt'
-    image_controlnet_path = 'models/image_controlnet.ckpt'
-    flow_controlnet_path = 'models/flow_controlnet.ckpt'
-    ImageConductor_net = ImageConductor(device=device, 
-                                        unet_path=unet_path, 
-                                        image_controlnet_path=image_controlnet_path, 
-                                        flow_controlnet_path=flow_controlnet_path, 
-                                        height=256,
-                                        width=384,
-                                        model_length=16
-                                        )
-    first_frame_path = gr.State()
-    tracking_points = gr.State([])
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            image_upload_button = gr.UploadButton(label="Upload Image",file_types=["image"])
-            add_drag_button = gr.Button(value="Add Drag")
-            reset_button = gr.Button(value="Reset")
-            delete_last_drag_button = gr.Button(value="Delete last drag")
-            delete_last_step_button = gr.Button(value="Delete last step")
-            
-            
-
-        with gr.Column(scale=7):
-            with gr.Row():
-                with gr.Column(scale=6):
-                    input_image = gr.Image(label="Input Image",
-                                        interactive=True,
-                                        height=265,
-                                        width=384,)
-                with gr.Column(scale=6):
-                    output_image = gr.Image(label="Motion Path",
-                                            interactive=False,
+    return {tracking_points_var: tracking_points, input_image: trajectory_map}
+
+
+if __name__=="__main__":
+    block = gr.Blocks(
+            theme=gr.themes.Soft(
+                radius_size=gr.themes.sizes.radius_none,
+                text_size=gr.themes.sizes.text_md
+            )
+            ).queue()
+    with block as demo:
+        with gr.Row():
+            with gr.Column():
+                gr.HTML(head)
+
+        gr.Markdown(descriptions)
+
+        with gr.Accordion(label="🛠️ Instructions:", open=True, elem_id="accordion"):
+            with gr.Row(equal_height=True):
+                gr.Markdown(instructions)      
+
+
+        # device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        device = torch.device("cuda")
+        unet_path = 'models/unet.ckpt'
+        image_controlnet_path = 'models/image_controlnet.ckpt'
+        flow_controlnet_path = 'models/flow_controlnet.ckpt'
+        ImageConductor_net = ImageConductor(device=device, 
+                                            unet_path=unet_path, 
+                                            image_controlnet_path=image_controlnet_path, 
+                                            flow_controlnet_path=flow_controlnet_path, 
                                             height=256,
-                                            width=384,)
-    with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(value="a wonderful elf.", label="Prompt (highly-recommended)", interactive=True, visible=True)
-            negative_prompt = gr.Text(
-                        label="Negative Prompt",
-                        max_lines=5,
-                        placeholder="Please input your negative prompt",
-                        value='worst quality, low quality, letterboxed',lines=1
-                    )
-            drag_mode = gr.Radio(['camera', 'object'], label='Drag mode: ', value='object', scale=2)
-            run_button = gr.Button(value="Run")
-
-            with gr.Accordion("More input params", open=False, elem_id="accordion1"):
-                with gr.Group():
-                    seed = gr.Textbox(
-                        label="Seed: ",  value=561793204,
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                                            width=384,
+                                            model_length=16
+                                            )
+        first_frame_path_var = gr.State(value=None)
+        tracking_points_var = gr.State([])
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_upload_button = gr.UploadButton(label="Upload Image",file_types=["image"])
+                add_drag_button = gr.Button(value="Add Drag")
+                reset_button = gr.Button(value="Reset")
+                delete_last_drag_button = gr.Button(value="Delete last drag")
+                delete_last_step_button = gr.Button(value="Delete last step")
+                
                 
-                with gr.Group():
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=1,
-                            maximum=12,
-                            step=0.1,
-                            value=8.5,
+
+            with gr.Column(scale=7):
+                with gr.Row():
+                    with gr.Column(scale=6):
+                        input_image = gr.Image(label="Input Image",
+                                            interactive=True,
+                                            height=300,
+                                            width=384,)
+                    with gr.Column(scale=6):
+                        output_image = gr.Image(label="Motion Path",
+                                                interactive=False,
+                                                height=256,
+                                                width=384,)
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt = gr.Textbox(value="a wonderful elf.", label="Prompt (highly-recommended)", interactive=True, visible=True)
+                negative_prompt = gr.Text(
+                            label="Negative Prompt",
+                            max_lines=5,
+                            placeholder="Please input your negative prompt",
+                            value='worst quality, low quality, letterboxed',lines=1
                         )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=25,
+                drag_mode = gr.Radio(['camera', 'object'], label='Drag mode: ', value='object', scale=2)
+                run_button = gr.Button(value="Run")
+
+                with gr.Accordion("More input params", open=False, elem_id="accordion1"):
+                    with gr.Group():
+                        seed = gr.Textbox(
+                            label="Seed: ",  value=561793204,
                         )
-                        
-                with gr.Group():
-                    personalized = gr.Dropdown(label="Personalized", choices=['HelloObject', 'TUSUN', ""], value="")
-                    examples_type = gr.Textbox(label="Examples Type (Ignore) ",  value="", visible=False)
-
-        with gr.Column(scale=7):
-            output_video = gr.Video(value=None, 
-                                    label="Output Video", 
-                                    width=384, 
-                                    height=256)
+                        randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                    
+                    with gr.Group():
+                        with gr.Row():
+                            guidance_scale = gr.Slider(
+                                label="Guidance scale",
+                                minimum=1,
+                                maximum=12,
+                                step=0.1,
+                                value=8.5,
+                            )
+                            num_inference_steps = gr.Slider(
+                                label="Number of inference steps",
+                                minimum=1,
+                                maximum=50,
+                                step=1,
+                                value=25,
+                            )
+                            
+                    with gr.Group():
+                        personalized = gr.Dropdown(label="Personalized", choices=['HelloObject', 'TUSUN', ""], value="")
+                        examples_type = gr.Textbox(label="Examples Type (Ignore) ",  value="", visible=False)
+
+            with gr.Column(scale=7):
+                output_video = gr.Video(
+                                        label="Output Video", 
+                                        width=384, 
+                                        height=256)
+                
+                
+        with gr.Row():
+            def process_example(input_image, prompt, drag_mode, seed, personalized, examples_type): 
+
+                return input_image, prompt, drag_mode, seed, personalized, examples_type
+    
+            example = gr.Examples(
+                label="Input Example",
+                examples=image_examples,
+                inputs=[input_image, prompt, drag_mode, seed, personalized, examples_type],
+                outputs=[input_image, prompt, drag_mode, seed, personalized, examples_type],
+                fn=process_example,
+                run_on_click=True,
+                examples_per_page=10,
+                cache_examples=False,
+            )
             
             
-    with gr.Row():
-        def process_example(input_image, prompt, drag_mode, seed, personalized, examples_type): 
-
-            return input_image, prompt, drag_mode, seed, personalized, examples_type
-  
-        example = gr.Examples(
-            label="Input Example",
-            examples=image_examples,
-            inputs=[input_image, prompt, drag_mode, seed, personalized, examples_type],
-            outputs=[input_image, prompt, drag_mode, seed, personalized, examples_type],
-            fn=process_example,
-            run_on_click=True,
-            examples_per_page=10,
-            cache_examples=False,
-        )
-        
-        
-    with gr.Row():
-        gr.Markdown(citation)
+        with gr.Row():
+            gr.Markdown(citation)
 
-    
-    image_upload_button.upload(preprocess_image, image_upload_button, [input_image, first_frame_path, tracking_points])
+        
+        image_upload_button.upload(preprocess_image, [image_upload_button, tracking_points_var], [input_image, first_frame_path_var, tracking_points_var])
 
-    add_drag_button.click(add_drag, tracking_points, tracking_points)
+        add_drag_button.click(add_drag, [tracking_points_var], tracking_points_var)
 
-    delete_last_drag_button.click(delete_last_drag, [tracking_points, first_frame_path, drag_mode], [tracking_points, input_image])
+        delete_last_drag_button.click(delete_last_drag, [tracking_points_var, first_frame_path_var, drag_mode], [tracking_points_var, input_image])
 
-    delete_last_step_button.click(delete_last_step, [tracking_points, first_frame_path, drag_mode], [tracking_points, input_image])
+        delete_last_step_button.click(delete_last_step, [tracking_points_var, first_frame_path_var, drag_mode], [tracking_points_var, input_image])
 
-    reset_button.click(reset_states, [first_frame_path, tracking_points], [input_image, first_frame_path, tracking_points])
+        reset_button.click(reset_states, [first_frame_path_var, tracking_points_var], [input_image, first_frame_path_var, tracking_points_var])
 
-    input_image.select(add_tracking_points, [tracking_points, first_frame_path, drag_mode], [tracking_points, input_image])
+        input_image.select(add_tracking_points, [tracking_points_var, first_frame_path_var, drag_mode], [tracking_points_var, input_image])
 
-    run_button.click(ImageConductor_net.run, [first_frame_path, tracking_points, prompt, drag_mode, 
-                                              negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, personalized, examples_type], 
-                                              [output_image, output_video])
+        run_button.click(ImageConductor_net.run, [first_frame_path_var, tracking_points_var, prompt, drag_mode, 
+                                                negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, personalized, examples_type], 
+                                                [output_image, output_video])
 
 demo.launch()

pipelines/pipeline_imagecoductor.py

@@ -404,7 +404,6 @@ class ImageConductorPipeline(DiffusionPipeline):
         obj_latents = copy.deepcopy(latents)
         cam_latents = copy.deepcopy(latents)
         
-        print("device", device)
         # Prepare extra step kwargs.
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
@@ -463,13 +462,7 @@ class ImageConductorPipeline(DiffusionPipeline):
                         controlnet_images_mask = controlnet_images_mask.half()
                         controlnet_flows = controlnet_flows.half()
                         text_embeddings = text_embeddings.half()
-                    print("controlnet_noisy_latents device", controlnet_noisy_latents.device)
-                    print("controlnet_prompt_embeds device", controlnet_prompt_embeds.device)
-                    print("controlnet_images device", controlnet_images.device)
-                    print("t", t.device)
-                    
-                    
-                    print("self.image_controlnet", self.image_controlnet.controlnet_mid_block.weight.device)
+                   
                     
                     img_down_block_additional_residuals, img_mid_block_additional_residuals = self.image_controlnet(
                                 controlnet_noisy_latents, t,