Spaces:
Paused
Paused
| import copy | |
| import os | |
| import shutil | |
| import cv2 | |
| import gradio as gr | |
| import modules.scripts as scripts | |
| from modules import images | |
| from modules.processing import process_images | |
| from modules.shared import opts | |
| from PIL import Image | |
| import numpy as np | |
| _BASEDIR = "/controlnet-m2m" | |
| _BASEFILE = "animation" | |
| def get_all_frames(video_path): | |
| if video_path is None: | |
| return None | |
| cap = cv2.VideoCapture(video_path) | |
| frame_list = [] | |
| if not cap.isOpened(): | |
| return | |
| while True: | |
| ret, frame = cap.read() | |
| if ret: | |
| frame_list.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) | |
| else: | |
| return frame_list | |
| def get_min_frame_num(video_list): | |
| min_frame_num = -1 | |
| for video in video_list: | |
| if video is None: | |
| continue | |
| else: | |
| frame_num = len(video) | |
| print(frame_num) | |
| if min_frame_num < 0: | |
| min_frame_num = frame_num | |
| elif frame_num < min_frame_num: | |
| min_frame_num = frame_num | |
| return min_frame_num | |
| def pil2cv(image): | |
| new_image = np.array(image, dtype=np.uint8) | |
| if new_image.ndim == 2: | |
| pass | |
| elif new_image.shape[2] == 3: | |
| new_image = new_image[:, :, ::-1] | |
| elif new_image.shape[2] == 4: | |
| new_image = new_image[:, :, [2, 1, 0, 3]] | |
| return new_image | |
| def save_gif(path, image_list, name, duration): | |
| tmp_dir = path + "/tmp/" | |
| if os.path.isdir(tmp_dir): | |
| shutil.rmtree(tmp_dir) | |
| os.mkdir(tmp_dir) | |
| for i, image in enumerate(image_list): | |
| images.save_image(image, tmp_dir, f"output_{i}") | |
| os.makedirs(f"{path}{_BASEDIR}", exist_ok=True) | |
| image_list[0].save(f"{path}{_BASEDIR}/{name}.gif", save_all=True, append_images=image_list[1:], optimize=False, duration=duration, loop=0) | |
| class Script(scripts.Script): | |
| def title(self): | |
| return "controlnet m2m" | |
| def show(self, is_img2img): | |
| return True | |
| def ui(self, is_img2img): | |
| # How the script's is displayed in the UI. See https://gradio.app/docs/#components | |
| # for the different UI components you can use and how to create them. | |
| # Most UI components can return a value, such as a boolean for a checkbox. | |
| # The returned values are passed to the run method as parameters. | |
| ctrls_group = () | |
| max_models = opts.data.get("control_net_unit_count", 3) | |
| with gr.Group(): | |
| with gr.Accordion("ControlNet-M2M", open = False): | |
| duration = gr.Slider(label=f"Duration", value=50.0, minimum=10.0, maximum=200.0, step=10, interactive=True, elem_id='controlnet_movie2movie_duration_slider') | |
| with gr.Tabs(): | |
| for i in range(max_models): | |
| with gr.Tab(f"ControlNet-{i}"): | |
| with gr.TabItem("Movie Input"): | |
| ctrls_group += (gr.Video(format='mp4', source='upload', elem_id = f"video_{i}"), ) | |
| with gr.TabItem("Image Input"): | |
| ctrls_group += (gr.Image(source='upload', brush_radius=20, mirror_webcam=False, type='numpy', tool='sketch', elem_id=f'image_{i}'), ) | |
| ctrls_group += (gr.Checkbox(label=f"Save preprocessed", value=False, elem_id = f"save_pre_{i}"),) | |
| ctrls_group += (duration,) | |
| return ctrls_group | |
| def run(self, p, *args): | |
| # This is where the additional processing is implemented. The parameters include | |
| # self, the model object "p" (a StableDiffusionProcessing class, see | |
| # processing.py), and the parameters returned by the ui method. | |
| # Custom functions can be defined here, and additional libraries can be imported | |
| # to be used in processing. The return value should be a Processed object, which is | |
| # what is returned by the process_images method. | |
| contents_num = opts.data.get("control_net_unit_count", 3) | |
| arg_num = 3 | |
| item_list = [] | |
| video_list = [] | |
| for input_set in [tuple(args[:contents_num * arg_num][i:i+3]) for i in range(0, len(args[:contents_num * arg_num]), arg_num)]: | |
| if input_set[0] is not None: | |
| item_list.append([get_all_frames(input_set[0]), "video"]) | |
| video_list.append(get_all_frames(input_set[0])) | |
| if input_set[1] is not None: | |
| item_list.append([cv2.cvtColor(pil2cv(input_set[1]["image"]), cv2.COLOR_BGRA2RGB), "image"]) | |
| save_pre = list(args[2:contents_num * arg_num:3]) | |
| item_num = len(item_list) | |
| video_num = len(video_list) | |
| duration, = args[contents_num * arg_num:] | |
| frame_num = get_min_frame_num(video_list) | |
| if frame_num > 0: | |
| output_image_list = [] | |
| pre_output_image_list = [] | |
| for i in range(item_num): | |
| pre_output_image_list.append([]) | |
| for frame in range(frame_num): | |
| copy_p = copy.copy(p) | |
| copy_p.control_net_input_image = [] | |
| for item in item_list: | |
| if item[1] == "video": | |
| copy_p.control_net_input_image.append(item[0][frame]) | |
| elif item[1] == "image": | |
| copy_p.control_net_input_image.append(item[0]) | |
| else: | |
| continue | |
| proc = process_images(copy_p) | |
| img = proc.images[0] | |
| output_image_list.append(img) | |
| for i in range(len(save_pre)): | |
| if save_pre[i]: | |
| try: | |
| pre_output_image_list[i].append(proc.images[i + 1]) | |
| except: | |
| print(f"proc.images[{i} failed") | |
| copy_p.close() | |
| # filename format is seq-seed-animation.gif seq is 5 places left filled with 0 | |
| seq = images.get_next_sequence_number(f"{p.outpath_samples}{_BASEDIR}", "") | |
| filename = f"{seq:05}-{proc.seed}-{_BASEFILE}" | |
| save_gif(p.outpath_samples, output_image_list, filename, duration) | |
| proc.images = [f"{p.outpath_samples}{_BASEDIR}/{filename}.gif"] | |
| for i in range(len(save_pre)): | |
| if save_pre[i]: | |
| # control files add -controlX.gif where X is the controlnet number | |
| save_gif(p.outpath_samples, pre_output_image_list[i], f"{filename}-control{i}", duration) | |
| proc.images.append(f"{p.outpath_samples}{_BASEDIR}/{filename}-control{i}.gif") | |
| else: | |
| proc = process_images(p) | |
| return proc | |