From f75196fd409c5190b38d50768eb3d9cfc11bfe8d Mon Sep 17 00:00:00 2001 From: julie wang Date: Wed, 27 Nov 2024 17:47:33 +0100 Subject: [PATCH] feat(ml): fix inference regeneration and crop canny --- scripts/gen_vid_diffusion.py | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/scripts/gen_vid_diffusion.py b/scripts/gen_vid_diffusion.py index de0732c3..f4706f87 100644 --- a/scripts/gen_vid_diffusion.py +++ b/scripts/gen_vid_diffusion.py @@ -610,18 +610,28 @@ def generate( # apply canny to the bbox to avoid edges around cond image img_tensor_canny = img_tensor[:, y0:y1, x0:x1] mask_canny = mask[:, y0:y1, x0:x1] + canny = fill_img_with_canny( + img_tensor_canny.unsqueeze(0), + mask_canny.unsqueeze(0), + low_threshold=alg_diffusion_sketch_canny_thresholds[0], + high_threshold=alg_diffusion_sketch_canny_thresholds[1], + low_threshold_random=alg_diffusion_sketch_canny_thresholds[0], + high_threshold_random=alg_diffusion_sketch_canny_thresholds[1], + select_canny=[select_canny_list[sequence_count]], + ) + else: img_tensor_canny = img_tensor mask_canny = mask - canny = fill_img_with_canny( - img_tensor_canny.unsqueeze(0), - mask_canny.unsqueeze(0), - low_threshold=alg_diffusion_sketch_canny_thresholds[0], - high_threshold=alg_diffusion_sketch_canny_thresholds[1], - low_threshold_random=alg_diffusion_sketch_canny_thresholds[0], - high_threshold_random=alg_diffusion_sketch_canny_thresholds[1], - select_canny=[select_canny_list[sequence_count]], - ) + canny = fill_img_with_canny( + img_tensor_canny.unsqueeze(0), + mask_canny.unsqueeze(0), + low_threshold=alg_diffusion_sketch_canny_thresholds[0], + high_threshold=alg_diffusion_sketch_canny_thresholds[1], + low_threshold_random=-1, + high_threshold_random=-1, + select_canny=[select_canny_list[sequence_count]], + ) if cond_in: # restore background cond_image = img_tensor.unsqueeze(0).clone()