167 | Apply ControlNet to a video 168 |
169 |You may also like:
181 |
229 | work with longer videos / skip the queue:
230 | """, elem_id="duplicate-container")
231 |
232 | inputs = [prompt,video_inp,control_task, seed_inp, trim_in, ddim_steps, scale]
233 | outputs = [video_out, share_group]
234 | #outputs = [status]
235 |
236 |
237 | gr.HTML(article)
238 |
239 | submit_btn.click(infer, inputs, outputs)
240 | share_button.click(None, [], [], _js=share_js)
241 |
242 |
243 |
244 | demo.launch().queue(max_size=12)
--------------------------------------------------------------------------------
/model.py:
--------------------------------------------------------------------------------
1 | # This file is adapted from gradio_*.py in https://github.com/lllyasviel/ControlNet/tree/f4748e3630d8141d7765e2bd9b1e348f47847707
2 | # The original license file is LICENSE.ControlNet in this repo.
3 | from __future__ import annotations
4 |
5 | import pathlib
6 | import random
7 | import shlex
8 | import subprocess
9 | import sys
10 |
11 | import cv2
12 | import einops
13 | import numpy as np
14 | import torch
15 | from pytorch_lightning import seed_everything
16 |
17 | sys.path.append('ControlNet')
18 |
19 | import config
20 | from annotator.canny import apply_canny
21 | from annotator.hed import apply_hed, nms
22 | from annotator.midas import apply_midas
23 | from annotator.mlsd import apply_mlsd
24 | from annotator.openpose import apply_openpose
25 | from annotator.uniformer import apply_uniformer
26 | from annotator.util import HWC3, resize_image
27 | from cldm.model import create_model, load_state_dict
28 | from ldm.models.diffusion.ddim import DDIMSampler
29 | from share import *
30 |
31 | ORIGINAL_MODEL_NAMES = {
32 | 'canny': 'control_sd15_canny.pth',
33 | 'hough': 'control_sd15_mlsd.pth',
34 | 'hed': 'control_sd15_hed.pth',
35 | 'scribble': 'control_sd15_scribble.pth',
36 | 'pose': 'control_sd15_openpose.pth',
37 | 'seg': 'control_sd15_seg.pth',
38 | 'depth': 'control_sd15_depth.pth',
39 | 'normal': 'control_sd15_normal.pth',
40 | }
41 | ORIGINAL_WEIGHT_ROOT = 'https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/'
42 |
43 | LIGHTWEIGHT_MODEL_NAMES = {
44 | 'canny': 'control_canny-fp16.safetensors',
45 | 'hough': 'control_mlsd-fp16.safetensors',
46 | 'hed': 'control_hed-fp16.safetensors',
47 | 'scribble': 'control_scribble-fp16.safetensors',
48 | 'pose': 'control_openpose-fp16.safetensors',
49 | 'seg': 'control_seg-fp16.safetensors',
50 | 'depth': 'control_depth-fp16.safetensors',
51 | 'normal': 'control_normal-fp16.safetensors',
52 | }
53 | LIGHTWEIGHT_WEIGHT_ROOT = 'https://huggingface.co/webui/ControlNet-modules-safetensors/resolve/main/'
54 |
55 |
56 | class Model:
57 | def __init__(self,
58 | model_config_path: str = 'ControlNet/models/cldm_v15.yaml',
59 | model_dir: str = 'models',
60 | use_lightweight: bool = True):
61 | self.device = torch.device(
62 | 'cuda:0' if torch.cuda.is_available() else 'cpu')
63 | self.model = create_model(model_config_path).to(self.device)
64 | self.ddim_sampler = DDIMSampler(self.model)
65 | self.task_name = ''
66 |
67 | self.model_dir = pathlib.Path(model_dir)
68 | self.model_dir.mkdir(exist_ok=True, parents=True)
69 |
70 | self.use_lightweight = use_lightweight
71 | if use_lightweight:
72 | self.model_names = LIGHTWEIGHT_MODEL_NAMES
73 | self.weight_root = LIGHTWEIGHT_WEIGHT_ROOT
74 | base_model_url = 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors'
75 | self.load_base_model(base_model_url)
76 | else:
77 | self.model_names = ORIGINAL_MODEL_NAMES
78 | self.weight_root = ORIGINAL_WEIGHT_ROOT
79 |
80 | self.download_models()
81 |
82 | def download_base_model(self, model_url: str) -> pathlib.Path:
83 | model_name = model_url.split('/')[-1]
84 | out_path = self.model_dir / model_name
85 | if not out_path.exists():
86 | subprocess.run(shlex.split(f'wget {model_url} -O {out_path}'))
87 | return out_path
88 |
89 | def load_base_model(self, model_url: str) -> None:
90 | model_path = self.download_base_model(model_url)
91 | self.model.load_state_dict(load_state_dict(model_path,
92 | location=self.device.type),
93 | strict=False)
94 |
95 | def load_weight(self, task_name: str) -> None:
96 | if task_name == self.task_name:
97 | return
98 | weight_path = self.get_weight_path(task_name)
99 | if not self.use_lightweight:
100 | self.model.load_state_dict(
101 | load_state_dict(weight_path, location=self.device))
102 | else:
103 | self.model.control_model.load_state_dict(
104 | load_state_dict(weight_path, location=self.device.type))
105 | self.task_name = task_name
106 |
107 | def get_weight_path(self, task_name: str) -> str:
108 | if 'scribble' in task_name:
109 | task_name = 'scribble'
110 | return f'{self.model_dir}/{self.model_names[task_name]}'
111 |
112 | def download_models(self) -> None:
113 | self.model_dir.mkdir(exist_ok=True, parents=True)
114 | for name in self.model_names.values():
115 | out_path = self.model_dir / name
116 | if out_path.exists():
117 | continue
118 | subprocess.run(
119 | shlex.split(f'wget {self.weight_root}{name} -O {out_path}'))
120 |
121 | @torch.inference_mode()
122 | def process_canny(self, input_image, prompt, a_prompt, n_prompt,
123 | num_samples, image_resolution, ddim_steps, scale, seed,
124 | eta, low_threshold, high_threshold):
125 | self.load_weight('canny')
126 |
127 | img = resize_image(HWC3(input_image), image_resolution)
128 | H, W, C = img.shape
129 |
130 | detected_map = apply_canny(img, low_threshold, high_threshold)
131 | detected_map = HWC3(detected_map)
132 |
133 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
134 | control = torch.stack([control for _ in range(num_samples)], dim=0)
135 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
136 |
137 | if seed == -1:
138 | seed = random.randint(0, 65535)
139 | seed_everything(seed)
140 |
141 | if config.save_memory:
142 | self.model.low_vram_shift(is_diffusing=False)
143 |
144 | cond = {
145 | 'c_concat': [control],
146 | 'c_crossattn': [
147 | self.model.get_learned_conditioning(
148 | [prompt + ', ' + a_prompt] * num_samples)
149 | ]
150 | }
151 | un_cond = {
152 | 'c_concat': [control],
153 | 'c_crossattn':
154 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
155 | }
156 | shape = (4, H // 8, W // 8)
157 |
158 | if config.save_memory:
159 | self.model.low_vram_shift(is_diffusing=True)
160 |
161 | samples, intermediates = self.ddim_sampler.sample(
162 | ddim_steps,
163 | num_samples,
164 | shape,
165 | cond,
166 | verbose=False,
167 | eta=eta,
168 | unconditional_guidance_scale=scale,
169 | unconditional_conditioning=un_cond)
170 |
171 | if config.save_memory:
172 | self.model.low_vram_shift(is_diffusing=False)
173 |
174 | x_samples = self.model.decode_first_stage(samples)
175 | x_samples = (
176 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
177 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
178 |
179 | results = [x_samples[i] for i in range(num_samples)]
180 | return [255 - detected_map] + results
181 |
182 | @torch.inference_mode()
183 | def process_hough(self, input_image, prompt, a_prompt, n_prompt,
184 | num_samples, image_resolution, detect_resolution,
185 | ddim_steps, scale, seed, eta, value_threshold,
186 | distance_threshold):
187 | self.load_weight('hough')
188 |
189 | input_image = HWC3(input_image)
190 | detected_map = apply_mlsd(resize_image(input_image, detect_resolution),
191 | value_threshold, distance_threshold)
192 | detected_map = HWC3(detected_map)
193 | img = resize_image(input_image, image_resolution)
194 | H, W, C = img.shape
195 |
196 | detected_map = cv2.resize(detected_map, (W, H),
197 | interpolation=cv2.INTER_NEAREST)
198 |
199 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
200 | control = torch.stack([control for _ in range(num_samples)], dim=0)
201 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
202 |
203 | if seed == -1:
204 | seed = random.randint(0, 65535)
205 | seed_everything(seed)
206 |
207 | if config.save_memory:
208 | self.model.low_vram_shift(is_diffusing=False)
209 |
210 | cond = {
211 | 'c_concat': [control],
212 | 'c_crossattn': [
213 | self.model.get_learned_conditioning(
214 | [prompt + ', ' + a_prompt] * num_samples)
215 | ]
216 | }
217 | un_cond = {
218 | 'c_concat': [control],
219 | 'c_crossattn':
220 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
221 | }
222 | shape = (4, H // 8, W // 8)
223 |
224 | if config.save_memory:
225 | self.model.low_vram_shift(is_diffusing=True)
226 |
227 | samples, intermediates = self.ddim_sampler.sample(
228 | ddim_steps,
229 | num_samples,
230 | shape,
231 | cond,
232 | verbose=False,
233 | eta=eta,
234 | unconditional_guidance_scale=scale,
235 | unconditional_conditioning=un_cond)
236 |
237 | if config.save_memory:
238 | self.model.low_vram_shift(is_diffusing=False)
239 |
240 | x_samples = self.model.decode_first_stage(samples)
241 | x_samples = (
242 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
243 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
244 |
245 | results = [x_samples[i] for i in range(num_samples)]
246 | return [
247 | 255 - cv2.dilate(detected_map,
248 | np.ones(shape=(3, 3), dtype=np.uint8),
249 | iterations=1)
250 | ] + results
251 |
252 | @torch.inference_mode()
253 | def process_hed(self, input_image, prompt, a_prompt, n_prompt, num_samples,
254 | image_resolution, detect_resolution, ddim_steps, scale,
255 | seed, eta):
256 | self.load_weight('hed')
257 |
258 | input_image = HWC3(input_image)
259 | detected_map = apply_hed(resize_image(input_image, detect_resolution))
260 | detected_map = HWC3(detected_map)
261 | img = resize_image(input_image, image_resolution)
262 | H, W, C = img.shape
263 |
264 | detected_map = cv2.resize(detected_map, (W, H),
265 | interpolation=cv2.INTER_LINEAR)
266 |
267 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
268 | control = torch.stack([control for _ in range(num_samples)], dim=0)
269 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
270 |
271 | if seed == -1:
272 | seed = random.randint(0, 65535)
273 | seed_everything(seed)
274 |
275 | if config.save_memory:
276 | self.model.low_vram_shift(is_diffusing=False)
277 |
278 | cond = {
279 | 'c_concat': [control],
280 | 'c_crossattn': [
281 | self.model.get_learned_conditioning(
282 | [prompt + ', ' + a_prompt] * num_samples)
283 | ]
284 | }
285 | un_cond = {
286 | 'c_concat': [control],
287 | 'c_crossattn':
288 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
289 | }
290 | shape = (4, H // 8, W // 8)
291 |
292 | if config.save_memory:
293 | self.model.low_vram_shift(is_diffusing=True)
294 |
295 | samples, intermediates = self.ddim_sampler.sample(
296 | ddim_steps,
297 | num_samples,
298 | shape,
299 | cond,
300 | verbose=False,
301 | eta=eta,
302 | unconditional_guidance_scale=scale,
303 | unconditional_conditioning=un_cond)
304 |
305 | if config.save_memory:
306 | self.model.low_vram_shift(is_diffusing=False)
307 |
308 | x_samples = self.model.decode_first_stage(samples)
309 | x_samples = (
310 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
311 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
312 |
313 | results = [x_samples[i] for i in range(num_samples)]
314 | return [detected_map] + results
315 |
316 | @torch.inference_mode()
317 | def process_scribble(self, input_image, prompt, a_prompt, n_prompt,
318 | num_samples, image_resolution, ddim_steps, scale,
319 | seed, eta):
320 | self.load_weight('scribble')
321 |
322 | img = resize_image(HWC3(input_image), image_resolution)
323 | H, W, C = img.shape
324 |
325 | detected_map = np.zeros_like(img, dtype=np.uint8)
326 | detected_map[np.min(img, axis=2) < 127] = 255
327 |
328 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
329 | control = torch.stack([control for _ in range(num_samples)], dim=0)
330 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
331 |
332 | if seed == -1:
333 | seed = random.randint(0, 65535)
334 | seed_everything(seed)
335 |
336 | if config.save_memory:
337 | self.model.low_vram_shift(is_diffusing=False)
338 |
339 | cond = {
340 | 'c_concat': [control],
341 | 'c_crossattn': [
342 | self.model.get_learned_conditioning(
343 | [prompt + ', ' + a_prompt] * num_samples)
344 | ]
345 | }
346 | un_cond = {
347 | 'c_concat': [control],
348 | 'c_crossattn':
349 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
350 | }
351 | shape = (4, H // 8, W // 8)
352 |
353 | if config.save_memory:
354 | self.model.low_vram_shift(is_diffusing=True)
355 |
356 | samples, intermediates = self.ddim_sampler.sample(
357 | ddim_steps,
358 | num_samples,
359 | shape,
360 | cond,
361 | verbose=False,
362 | eta=eta,
363 | unconditional_guidance_scale=scale,
364 | unconditional_conditioning=un_cond)
365 |
366 | if config.save_memory:
367 | self.model.low_vram_shift(is_diffusing=False)
368 |
369 | x_samples = self.model.decode_first_stage(samples)
370 | x_samples = (
371 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
372 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
373 |
374 | results = [x_samples[i] for i in range(num_samples)]
375 | return [255 - detected_map] + results
376 |
377 | @torch.inference_mode()
378 | def process_scribble_interactive(self, input_image, prompt, a_prompt,
379 | n_prompt, num_samples, image_resolution,
380 | ddim_steps, scale, seed, eta):
381 | self.load_weight('scribble')
382 |
383 | img = resize_image(HWC3(input_image['mask'][:, :, 0]),
384 | image_resolution)
385 | H, W, C = img.shape
386 |
387 | detected_map = np.zeros_like(img, dtype=np.uint8)
388 | detected_map[np.min(img, axis=2) > 127] = 255
389 |
390 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
391 | control = torch.stack([control for _ in range(num_samples)], dim=0)
392 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
393 |
394 | if seed == -1:
395 | seed = random.randint(0, 65535)
396 | seed_everything(seed)
397 |
398 | if config.save_memory:
399 | self.model.low_vram_shift(is_diffusing=False)
400 |
401 | cond = {
402 | 'c_concat': [control],
403 | 'c_crossattn': [
404 | self.model.get_learned_conditioning(
405 | [prompt + ', ' + a_prompt] * num_samples)
406 | ]
407 | }
408 | un_cond = {
409 | 'c_concat': [control],
410 | 'c_crossattn':
411 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
412 | }
413 | shape = (4, H // 8, W // 8)
414 |
415 | if config.save_memory:
416 | self.model.low_vram_shift(is_diffusing=True)
417 |
418 | samples, intermediates = self.ddim_sampler.sample(
419 | ddim_steps,
420 | num_samples,
421 | shape,
422 | cond,
423 | verbose=False,
424 | eta=eta,
425 | unconditional_guidance_scale=scale,
426 | unconditional_conditioning=un_cond)
427 |
428 | if config.save_memory:
429 | self.model.low_vram_shift(is_diffusing=False)
430 |
431 | x_samples = self.model.decode_first_stage(samples)
432 | x_samples = (
433 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
434 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
435 |
436 | results = [x_samples[i] for i in range(num_samples)]
437 | return [255 - detected_map] + results
438 |
439 | @torch.inference_mode()
440 | def process_fake_scribble(self, input_image, prompt, a_prompt, n_prompt,
441 | num_samples, image_resolution, detect_resolution,
442 | ddim_steps, scale, seed, eta):
443 | self.load_weight('scribble')
444 |
445 | input_image = HWC3(input_image)
446 | detected_map = apply_hed(resize_image(input_image, detect_resolution))
447 | detected_map = HWC3(detected_map)
448 | img = resize_image(input_image, image_resolution)
449 | H, W, C = img.shape
450 |
451 | detected_map = cv2.resize(detected_map, (W, H),
452 | interpolation=cv2.INTER_LINEAR)
453 | detected_map = nms(detected_map, 127, 3.0)
454 | detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
455 | detected_map[detected_map > 4] = 255
456 | detected_map[detected_map < 255] = 0
457 |
458 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
459 | control = torch.stack([control for _ in range(num_samples)], dim=0)
460 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
461 |
462 | if seed == -1:
463 | seed = random.randint(0, 65535)
464 | seed_everything(seed)
465 |
466 | if config.save_memory:
467 | self.model.low_vram_shift(is_diffusing=False)
468 |
469 | cond = {
470 | 'c_concat': [control],
471 | 'c_crossattn': [
472 | self.model.get_learned_conditioning(
473 | [prompt + ', ' + a_prompt] * num_samples)
474 | ]
475 | }
476 | un_cond = {
477 | 'c_concat': [control],
478 | 'c_crossattn':
479 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
480 | }
481 | shape = (4, H // 8, W // 8)
482 |
483 | if config.save_memory:
484 | self.model.low_vram_shift(is_diffusing=True)
485 |
486 | samples, intermediates = self.ddim_sampler.sample(
487 | ddim_steps,
488 | num_samples,
489 | shape,
490 | cond,
491 | verbose=False,
492 | eta=eta,
493 | unconditional_guidance_scale=scale,
494 | unconditional_conditioning=un_cond)
495 |
496 | if config.save_memory:
497 | self.model.low_vram_shift(is_diffusing=False)
498 |
499 | x_samples = self.model.decode_first_stage(samples)
500 | x_samples = (
501 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
502 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
503 |
504 | results = [x_samples[i] for i in range(num_samples)]
505 | return [255 - detected_map] + results
506 |
507 | @torch.inference_mode()
508 | def process_pose(self, input_image, prompt, a_prompt, n_prompt,
509 | num_samples, image_resolution, detect_resolution,
510 | ddim_steps, scale, seed, eta):
511 | self.load_weight('pose')
512 |
513 | input_image = HWC3(input_image)
514 | detected_map, _ = apply_openpose(
515 | resize_image(input_image, detect_resolution))
516 | detected_map = HWC3(detected_map)
517 | img = resize_image(input_image, image_resolution)
518 | H, W, C = img.shape
519 |
520 | detected_map = cv2.resize(detected_map, (W, H),
521 | interpolation=cv2.INTER_NEAREST)
522 |
523 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
524 | control = torch.stack([control for _ in range(num_samples)], dim=0)
525 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
526 |
527 | if seed == -1:
528 | seed = random.randint(0, 65535)
529 | seed_everything(seed)
530 |
531 | if config.save_memory:
532 | self.model.low_vram_shift(is_diffusing=False)
533 |
534 | cond = {
535 | 'c_concat': [control],
536 | 'c_crossattn': [
537 | self.model.get_learned_conditioning(
538 | [prompt + ', ' + a_prompt] * num_samples)
539 | ]
540 | }
541 | un_cond = {
542 | 'c_concat': [control],
543 | 'c_crossattn':
544 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
545 | }
546 | shape = (4, H // 8, W // 8)
547 |
548 | if config.save_memory:
549 | self.model.low_vram_shift(is_diffusing=True)
550 |
551 | samples, intermediates = self.ddim_sampler.sample(
552 | ddim_steps,
553 | num_samples,
554 | shape,
555 | cond,
556 | verbose=False,
557 | eta=eta,
558 | unconditional_guidance_scale=scale,
559 | unconditional_conditioning=un_cond)
560 |
561 | if config.save_memory:
562 | self.model.low_vram_shift(is_diffusing=False)
563 |
564 | x_samples = self.model.decode_first_stage(samples)
565 | x_samples = (
566 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
567 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
568 |
569 | results = [x_samples[i] for i in range(num_samples)]
570 | return [detected_map] + results
571 |
572 | @torch.inference_mode()
573 | def process_seg(self, input_image, prompt, a_prompt, n_prompt, num_samples,
574 | image_resolution, detect_resolution, ddim_steps, scale,
575 | seed, eta):
576 | self.load_weight('seg')
577 |
578 | input_image = HWC3(input_image)
579 | detected_map = apply_uniformer(
580 | resize_image(input_image, detect_resolution))
581 | img = resize_image(input_image, image_resolution)
582 | H, W, C = img.shape
583 |
584 | detected_map = cv2.resize(detected_map, (W, H),
585 | interpolation=cv2.INTER_NEAREST)
586 |
587 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
588 | control = torch.stack([control for _ in range(num_samples)], dim=0)
589 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
590 |
591 | if seed == -1:
592 | seed = random.randint(0, 65535)
593 | seed_everything(seed)
594 |
595 | if config.save_memory:
596 | self.model.low_vram_shift(is_diffusing=False)
597 |
598 | cond = {
599 | 'c_concat': [control],
600 | 'c_crossattn': [
601 | self.model.get_learned_conditioning(
602 | [prompt + ', ' + a_prompt] * num_samples)
603 | ]
604 | }
605 | un_cond = {
606 | 'c_concat': [control],
607 | 'c_crossattn':
608 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
609 | }
610 | shape = (4, H // 8, W // 8)
611 |
612 | if config.save_memory:
613 | self.model.low_vram_shift(is_diffusing=True)
614 |
615 | samples, intermediates = self.ddim_sampler.sample(
616 | ddim_steps,
617 | num_samples,
618 | shape,
619 | cond,
620 | verbose=False,
621 | eta=eta,
622 | unconditional_guidance_scale=scale,
623 | unconditional_conditioning=un_cond)
624 |
625 | if config.save_memory:
626 | self.model.low_vram_shift(is_diffusing=False)
627 |
628 | x_samples = self.model.decode_first_stage(samples)
629 | x_samples = (
630 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
631 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
632 |
633 | results = [x_samples[i] for i in range(num_samples)]
634 | return [detected_map] + results
635 |
636 | @torch.inference_mode()
637 | def process_depth(self, input_image, prompt, a_prompt, n_prompt,
638 | num_samples, image_resolution, detect_resolution,
639 | ddim_steps, scale, seed, eta):
640 | self.load_weight('depth')
641 |
642 | input_image = HWC3(input_image)
643 | detected_map, _ = apply_midas(
644 | resize_image(input_image, detect_resolution))
645 | detected_map = HWC3(detected_map)
646 | img = resize_image(input_image, image_resolution)
647 | H, W, C = img.shape
648 |
649 | detected_map = cv2.resize(detected_map, (W, H),
650 | interpolation=cv2.INTER_LINEAR)
651 |
652 | control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
653 | control = torch.stack([control for _ in range(num_samples)], dim=0)
654 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
655 |
656 | if seed == -1:
657 | seed = random.randint(0, 65535)
658 | seed_everything(seed)
659 |
660 | if config.save_memory:
661 | self.model.low_vram_shift(is_diffusing=False)
662 |
663 | cond = {
664 | 'c_concat': [control],
665 | 'c_crossattn': [
666 | self.model.get_learned_conditioning(
667 | [prompt + ', ' + a_prompt] * num_samples)
668 | ]
669 | }
670 | un_cond = {
671 | 'c_concat': [control],
672 | 'c_crossattn':
673 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
674 | }
675 | shape = (4, H // 8, W // 8)
676 |
677 | if config.save_memory:
678 | self.model.low_vram_shift(is_diffusing=True)
679 |
680 | samples, intermediates = self.ddim_sampler.sample(
681 | ddim_steps,
682 | num_samples,
683 | shape,
684 | cond,
685 | verbose=False,
686 | eta=eta,
687 | unconditional_guidance_scale=scale,
688 | unconditional_conditioning=un_cond)
689 |
690 | if config.save_memory:
691 | self.model.low_vram_shift(is_diffusing=False)
692 |
693 | x_samples = self.model.decode_first_stage(samples)
694 | x_samples = (
695 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
696 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
697 |
698 | results = [x_samples[i] for i in range(num_samples)]
699 | return [detected_map] + results
700 |
701 | @torch.inference_mode()
702 | def process_normal(self, input_image, prompt, a_prompt, n_prompt,
703 | num_samples, image_resolution, detect_resolution,
704 | ddim_steps, scale, seed, eta, bg_threshold):
705 | self.load_weight('normal')
706 |
707 | input_image = HWC3(input_image)
708 | _, detected_map = apply_midas(resize_image(input_image,
709 | detect_resolution),
710 | bg_th=bg_threshold)
711 | detected_map = HWC3(detected_map)
712 | img = resize_image(input_image, image_resolution)
713 | H, W, C = img.shape
714 |
715 | detected_map = cv2.resize(detected_map, (W, H),
716 | interpolation=cv2.INTER_LINEAR)
717 |
718 | control = torch.from_numpy(
719 | detected_map[:, :, ::-1].copy()).float().cuda() / 255.0
720 | control = torch.stack([control for _ in range(num_samples)], dim=0)
721 | control = einops.rearrange(control, 'b h w c -> b c h w').clone()
722 |
723 | if seed == -1:
724 | seed = random.randint(0, 65535)
725 | seed_everything(seed)
726 |
727 | if config.save_memory:
728 | self.model.low_vram_shift(is_diffusing=False)
729 |
730 | cond = {
731 | 'c_concat': [control],
732 | 'c_crossattn': [
733 | self.model.get_learned_conditioning(
734 | [prompt + ', ' + a_prompt] * num_samples)
735 | ]
736 | }
737 | un_cond = {
738 | 'c_concat': [control],
739 | 'c_crossattn':
740 | [self.model.get_learned_conditioning([n_prompt] * num_samples)]
741 | }
742 | shape = (4, H // 8, W // 8)
743 |
744 | if config.save_memory:
745 | self.model.low_vram_shift(is_diffusing=True)
746 |
747 | samples, intermediates = self.ddim_sampler.sample(
748 | ddim_steps,
749 | num_samples,
750 | shape,
751 | cond,
752 | verbose=False,
753 | eta=eta,
754 | unconditional_guidance_scale=scale,
755 | unconditional_conditioning=un_cond)
756 |
757 | if config.save_memory:
758 | self.model.low_vram_shift(is_diffusing=False)
759 |
760 | x_samples = self.model.decode_first_stage(samples)
761 | x_samples = (
762 | einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
763 | 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
764 |
765 | results = [x_samples[i] for i in range(num_samples)]
766 | return [detected_map] + results
767 |
--------------------------------------------------------------------------------
/patch:
--------------------------------------------------------------------------------
1 | diff --git a/annotator/hed/__init__.py b/annotator/hed/__init__.py
2 | index 42d8dc6..1587035 100644
3 | --- a/annotator/hed/__init__.py
4 | +++ b/annotator/hed/__init__.py
5 | @@ -1,8 +1,12 @@
6 | +import pathlib
7 | +
8 | import numpy as np
9 | import cv2
10 | import torch
11 | from einops import rearrange
12 |
13 | +root_dir = pathlib.Path(__file__).parents[2]
14 | +
15 |
16 | class Network(torch.nn.Module):
17 | def __init__(self):
18 | @@ -64,7 +68,7 @@ class Network(torch.nn.Module):
19 | torch.nn.Sigmoid()
20 | )
21 |
22 | - self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load('./annotator/ckpts/network-bsds500.pth').items()})
23 | + self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load(f'{root_dir}/annotator/ckpts/network-bsds500.pth').items()})
24 | # end
25 |
26 | def forward(self, tenInput):
27 | diff --git a/annotator/midas/api.py b/annotator/midas/api.py
28 | index 9fa305e..d8594ea 100644
29 | --- a/annotator/midas/api.py
30 | +++ b/annotator/midas/api.py
31 | @@ -1,5 +1,7 @@
32 | # based on https://github.com/isl-org/MiDaS
33 |
34 | +import pathlib
35 | +
36 | import cv2
37 | import torch
38 | import torch.nn as nn
39 | @@ -10,10 +12,11 @@ from .midas.midas_net import MidasNet
40 | from .midas.midas_net_custom import MidasNet_small
41 | from .midas.transforms import Resize, NormalizeImage, PrepareForNet
42 |
43 | +root_dir = pathlib.Path(__file__).parents[2]
44 |
45 | ISL_PATHS = {
46 | - "dpt_large": "annotator/ckpts/dpt_large-midas-2f21e586.pt",
47 | - "dpt_hybrid": "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",
48 | + "dpt_large": f"{root_dir}/annotator/ckpts/dpt_large-midas-2f21e586.pt",
49 | + "dpt_hybrid": f"{root_dir}/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",
50 | "midas_v21": "",
51 | "midas_v21_small": "",
52 | }
53 | diff --git a/annotator/mlsd/__init__.py b/annotator/mlsd/__init__.py
54 | index 75db717..f310fe6 100644
55 | --- a/annotator/mlsd/__init__.py
56 | +++ b/annotator/mlsd/__init__.py
57 | @@ -1,3 +1,5 @@
58 | +import pathlib
59 | +
60 | import cv2
61 | import numpy as np
62 | import torch
63 | @@ -8,8 +10,9 @@ from .models.mbv2_mlsd_tiny import MobileV2_MLSD_Tiny
64 | from .models.mbv2_mlsd_large import MobileV2_MLSD_Large
65 | from .utils import pred_lines
66 |
67 | +root_dir = pathlib.Path(__file__).parents[2]
68 |
69 | -model_path = './annotator/ckpts/mlsd_large_512_fp32.pth'
70 | +model_path = f'{root_dir}/annotator/ckpts/mlsd_large_512_fp32.pth'
71 | model = MobileV2_MLSD_Large()
72 | model.load_state_dict(torch.load(model_path), strict=True)
73 | model = model.cuda().eval()
74 | diff --git a/annotator/openpose/__init__.py b/annotator/openpose/__init__.py
75 | index 47d50a5..2369eed 100644
76 | --- a/annotator/openpose/__init__.py
77 | +++ b/annotator/openpose/__init__.py
78 | @@ -1,4 +1,5 @@
79 | import os
80 | +import pathlib
81 | os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
82 |
83 | import torch
84 | @@ -7,8 +8,10 @@ from . import util
85 | from .body import Body
86 | from .hand import Hand
87 |
88 | -body_estimation = Body('./annotator/ckpts/body_pose_model.pth')
89 | -hand_estimation = Hand('./annotator/ckpts/hand_pose_model.pth')
90 | +root_dir = pathlib.Path(__file__).parents[2]
91 | +
92 | +body_estimation = Body(f'{root_dir}/annotator/ckpts/body_pose_model.pth')
93 | +hand_estimation = Hand(f'{root_dir}/annotator/ckpts/hand_pose_model.pth')
94 |
95 |
96 | def apply_openpose(oriImg, hand=False):
97 | diff --git a/annotator/uniformer/__init__.py b/annotator/uniformer/__init__.py
98 | index 500e53c..4061dbe 100644
99 | --- a/annotator/uniformer/__init__.py
100 | +++ b/annotator/uniformer/__init__.py
101 | @@ -1,9 +1,12 @@
102 | +import pathlib
103 | +
104 | from annotator.uniformer.mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot
105 | from annotator.uniformer.mmseg.core.evaluation import get_palette
106 |
107 | +root_dir = pathlib.Path(__file__).parents[2]
108 |
109 | -checkpoint_file = "annotator/ckpts/upernet_global_small.pth"
110 | -config_file = 'annotator/uniformer/exp/upernet_global_small/config.py'
111 | +checkpoint_file = f"{root_dir}/annotator/ckpts/upernet_global_small.pth"
112 | +config_file = f'{root_dir}/annotator/uniformer/exp/upernet_global_small/config.py'
113 | model = init_segmentor(config_file, checkpoint_file).cuda()
114 |
115 |
116 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | opencv-python
2 | ffmpeg-python
3 | moviepy
4 | addict==2.4.0
5 | albumentations==1.3.0
6 | einops==0.6.0
7 | gradio==3.18.0
8 | imageio==2.25.0
9 | imageio-ffmpeg==0.4.8
10 | kornia==0.6.9
11 | omegaconf==2.3.0
12 | open-clip-torch==2.13.0
13 | opencv-contrib-python==4.7.0.68
14 | opencv-python-headless==4.7.0.68
15 | prettytable==3.6.0
16 | pytorch-lightning==1.9.0
17 | safetensors==0.2.8
18 | timm==0.6.12
19 | torch==1.13.1
20 | torchvision==0.14.1
21 | transformers==4.26.1
22 | xformers==0.0.16
23 | yapf==0.32.0
24 |
--------------------------------------------------------------------------------
/share_btn.py:
--------------------------------------------------------------------------------
1 | community_icon_html = """"""
5 |
6 | loading_icon_html = """"""
10 |
11 | share_js = """async () => {
12 | async function uploadFile(file){
13 | const UPLOAD_URL = 'https://huggingface.co/uploads';
14 | const response = await fetch(UPLOAD_URL, {
15 | method: 'POST',
16 | headers: {
17 | 'Content-Type': file.type,
18 | 'X-Requested-With': 'XMLHttpRequest',
19 | },
20 | body: file, /// <- File inherits from Blob
21 | });
22 | const url = await response.text();
23 | return url;
24 | }
25 |
26 | async function getVideoBlobFile(videoEL){
27 | const res = await fetch(videoEL.src);
28 | const blob = await res.blob();
29 | const videoId = Date.now() % 200;
30 | const fileName = `vid-pix2pix-${{videoId}}.wav`;
31 | const videoBlob = new File([blob], fileName, { type: 'video/mp4' });
32 | console.log(videoBlob);
33 | return videoBlob;
34 | }
35 |
36 | const gradioEl = document.querySelector("gradio-app").shadowRoot || document.querySelector('body > gradio-app');
37 | const captionTxt = gradioEl.querySelector('#prompt-in textarea').value;
38 | const inputVidEl = gradioEl.querySelector('#input-vid video');
39 | const outputVideo = gradioEl.querySelector('#video-output video');
40 |
41 | const shareBtnEl = gradioEl.querySelector('#share-btn');
42 | const shareIconEl = gradioEl.querySelector('#share-btn-share-icon');
43 | const loadingIconEl = gradioEl.querySelector('#share-btn-loading-icon');
44 | if(!outputVideo){
45 | return;
46 | };
47 | shareBtnEl.style.pointerEvents = 'none';
48 | shareIconEl.style.display = 'none';
49 | loadingIconEl.style.removeProperty('display');
50 |
51 | const inputFile = await getVideoBlobFile(inputVidEl);
52 | const urlInputVid = await uploadFile(inputFile);
53 | const videoOutFile = await getVideoBlobFile(outputVideo);
54 | const dataOutputVid = await uploadFile(videoOutFile);
55 |
56 | const descriptionMd = `
57 | #### Video input:
58 | ${urlInputVid}
59 |
60 | #### ControlNet result:
61 | ${dataOutputVid}
62 | `;
63 | const params = new URLSearchParams({
64 | title: captionTxt,
65 | description: descriptionMd,
66 | });
67 | const paramsStr = params.toString();
68 | window.open(`https://huggingface.co/spaces/fffiloni/ControlNet-Video/discussions/new?${paramsStr}`, '_blank');
69 | shareBtnEl.style.removeProperty('pointer-events');
70 | shareIconEl.style.removeProperty('display');
71 | loadingIconEl.style.display = 'none';
72 | }"""
--------------------------------------------------------------------------------
/style.css:
--------------------------------------------------------------------------------
1 | #col-container {max-width: 820px; margin-left: auto; margin-right: auto;}
2 | #duplicate-container{
3 | display: flex;
4 | justify-content: space-between;
5 | align-items: center;
6 | line-height: 1em;
7 | flex-direction: row-reverse;
8 | font-size:1em;
9 | }
10 | a, a:hover, a:visited {
11 | text-decoration-line: underline;
12 | font-weight: 600;
13 | color: #1f2937 !important;
14 | }
15 |
16 | .dark a, .dark a:hover, .dark a:visited {
17 | color: #f3f4f6 !important;
18 | }
19 |
20 | .footer {
21 | margin-bottom: 45px;
22 | margin-top: 10px;
23 | text-align: center;
24 | border-bottom: 1px solid #e5e5e5;
25 | }
26 |
27 | .footer>p {
28 | font-size: .8rem!important;
29 | display: inline-block;
30 | padding: 0 10px;
31 | transform: translateY(26px);
32 | background: white;
33 | }
34 | .dark .footer {
35 | border-color: #303030;
36 | }
37 | .dark .footer>p {
38 | background: #0b0f19;
39 | }
40 |
41 | div#may-like-container > p {
42 | font-size: .8em;
43 | margin-bottom: 4px;
44 | }
45 |
46 | .animate-spin {
47 | animation: spin 1s linear infinite;
48 | }
49 |
50 | @keyframes spin {
51 | from {
52 | transform: rotate(0deg);
53 | }
54 | to {
55 | transform: rotate(360deg);
56 | }
57 | }
58 |
59 | #share-btn-container {
60 | display: flex;
61 | padding-left: 0.5rem !important;
62 | padding-right: 0.5rem !important;
63 | background-color: #000000;
64 | justify-content: center;
65 | align-items: center;
66 | border-radius: 9999px !important;
67 | max-width: 13rem;
68 | }
69 |
70 | #share-btn-container:hover {
71 | background-color: #060606;
72 | }
73 |
74 | #share-btn {
75 | all: initial;
76 | color: #ffffff;
77 | font-weight: 600;
78 | cursor:pointer;
79 | font-family: 'IBM Plex Sans', sans-serif;
80 | margin-left: 0.5rem !important;
81 | padding-top: 0.5rem !important;
82 | padding-bottom: 0.5rem !important;
83 | right:0;
84 | }
85 |
86 | #share-btn * {
87 | all: unset;
88 | }
89 |
90 | #share-btn-container div:nth-child(-n+2){
91 | width: auto !important;
92 | min-height: 0px !important;
93 | }
94 |
95 | #share-btn-container .wrap {
96 | display: none !important;
97 | }
98 |
99 | #share-btn-container.hidden {
100 | display: none!important;
101 | }
--------------------------------------------------------------------------------