envs

app.py

@@ -302,30 +302,26 @@ class ImageConductor:
         self.blur_kernel = blur_kernel
 
     @spaces.GPU(duration=120)
-    def run(self, first_frame_path, tracking_points, prompt, drag_mode, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, personalized, examples_type):
-        print("Run!")
-        # if examples_type != "":
-        #     ### for adapting high version gradio
-        #     tracking_points = gr.State([])
-        #     first_frame_path = IMAGE_PATH[examples_type]
-        #     points = json.load(open(POINTS[examples_type]))
-        #     tracking_points.value.extend(points)
-        #     print("example first_frame_path", first_frame_path)
-        #     print("example tracking_points", tracking_points.value)
-            
+    def run(self, first_frame_path, tracking_points, prompt, drag_mode, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, personalized,):
+      
+        
         original_width, original_height=384, 256
         if isinstance(tracking_points, list):
             input_all_points = tracking_points
         else:
             input_all_points = tracking_points.value
         
-        print("input_all_points", input_all_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]
-        print("first_frame_path", first_frame_path)
+
         dir, base, ext = split_filename(first_frame_path)
         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        
@@ -337,8 +333,9 @@ class ImageConductor:
                 transforms.ToTensor(),
             ])
 
+        image_norm = lambda x: x
         image_paths = [first_frame_path]
-        controlnet_images = [(image_transforms(Image.open(path).convert("RGB"))) for path in image_paths]
+        controlnet_images = [image_norm(image_transforms(Image.open(path).convert("RGB"))) for path in image_paths]
         controlnet_images = torch.stack(controlnet_images).unsqueeze(0).to(device)
         controlnet_images = rearrange(controlnet_images, "b f c h w -> b c f h w")
         num_controlnet_images = controlnet_images.shape[2]
@@ -398,9 +395,10 @@ 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'})
         
+        
         outputs_path = os.path.join(output_dir, f'output_{i}_{id}.gif')
         save_videos_grid(sample[0][None], outputs_path)
-        print("Done!")
+
         return {output_image: visualized_drag, output_video: outputs_path}
 
 
@@ -410,7 +408,7 @@ def reset_states(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, tracking_points):
+def preprocess_image(image):
     image_pil = image2pil(image.name)
     raw_w, raw_h = image_pil.size
     resize_ratio = max(384/raw_w, 256/raw_h)
@@ -419,8 +417,7 @@ def preprocess_image(image, tracking_points):
     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)
-    tracking_points =  gr.State([])
-    return {input_image: first_frame_path, first_frame_path_var: first_frame_path, tracking_points_var: tracking_points, personalized:""}
+    return {input_image: first_frame_path, first_frame_path_var: first_frame_path, tracking_points_var: gr.State([]), personalized: ""}
 
 
 def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.SelectData):  # SelectData is a subclass of EventData
@@ -429,27 +426,14 @@ def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.Se
     elif drag_mode=='camera':
         color = (0, 0, 255, 255)
 
-    if not isinstance(tracking_points ,list):
-        print(f"You selected {evt.value} at {evt.index} from {evt.target}")
-        tracking_points.value[-1].append(evt.index)
-        print(tracking_points.value)
-        tracking_points_values =  tracking_points.value
-    else:
-        try:
-            tracking_points[-1].append(evt.index)
-        except Exception as e:
-            tracking_points.append([])
-            tracking_points[-1].append(evt.index)
-            print(f"Solved Error: {e}")
-        
-        tracking_points_values = tracking_points
-        
+    print(f"You selected {evt.value} at {evt.index} from {evt.target}")
+    tracking_points.value[-1].append(evt.index)
+    print(tracking_points.value)
     
     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_values:
+    for track in tracking_points.value:
         if len(track) > 1:
             for i in range(len(track)-1):
                 start_point = track[i]
@@ -470,12 +454,9 @@ def add_tracking_points(tracking_points, first_frame_path, drag_mode, evt: gr.Se
 
 
 def add_drag(tracking_points):
-    if not isinstance(tracking_points ,list):
-        # print("before", tracking_points.value)
-        tracking_points.value.append([])
-        # print(tracking_points.value)
-    else:
-        tracking_points.append([])
+    # import ipdb; ipdb.set_trace()
+    tracking_points.value.append([])
+    print(tracking_points.value)
     return {tracking_points_var: tracking_points}
     
 
@@ -537,142 +518,144 @@ def delete_last_step(tracking_points, first_frame_path, drag_mode):
     return {tracking_points_var: tracking_points, input_image: trajectory_map}
 
 
-block = gr.Blocks(
-        theme=gr.themes.Soft(
-            radius_size=gr.themes.sizes.radius_none,
-            text_size=gr.themes.sizes.text_md
-        )
-        )
-with block:
-    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_var = gr.State()
-    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.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,
+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(
-            #                         label="Output Video", 
-            #                         width=384, 
-            #                         height=256)
-            output_video = gr.Image(label="Output Video",
+                        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)
+                output_video = gr.Image(label="Output Video",
                                             height=256,
                                             width=384,)
-            
-            
-    with gr.Row():
-   
-
-        example = gr.Examples(
+                
+                
+        with gr.Row():
+            def process_examples(input_image, prompt, drag_mode, seed, personalized, first_frame_path_var, tracking_points_var):
+                return input_image, prompt, drag_mode, seed, personalized, first_frame_path_var, tracking_points_var
+            example = gr.Examples(
             label="Input Example",
             examples=image_examples,
             inputs=[input_image, prompt, drag_mode, seed, personalized, first_frame_path_var, tracking_points_var],
+            outputs=[input_image, prompt, drag_mode, seed, personalized, first_frame_path_var, tracking_points_var],
+            fn=process_examples,
             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, tracking_points_var], [input_image, first_frame_path_var, tracking_points_var, personalized])
+        
+        image_upload_button.upload(preprocess_image, image_upload_button, [input_image, first_frame_path_var, tracking_points_var, personalized])
 
-    add_drag_button.click(add_drag, tracking_points_var, tracking_points_var)
+        add_drag_button.click(add_drag, [tracking_points_var], tracking_points_var)
 
-    delete_last_drag_button.click(delete_last_drag, [tracking_points_var, first_frame_path_var, drag_mode], [tracking_points_var, 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_var, first_frame_path_var, drag_mode], [tracking_points_var, 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_var, tracking_points_var], [input_image, first_frame_path_var, tracking_points_var])
+        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_var, first_frame_path_var, drag_mode], [tracking_points_var, 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_var, tracking_points_var, 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], 
+                                                [output_image, output_video])
 
 block.queue().launch()