add video UI

This commit is contained in:
Artiprocher
2024-01-09 19:07:16 +08:00
parent 49db42bf1f
commit cd6fa30c9a
9 changed files with 294 additions and 20 deletions

View File

@@ -23,13 +23,11 @@ class MultiControlNetManager:
self.models = [unit.model for unit in controlnet_units]
self.scales = [unit.scale for unit in controlnet_units]
def process_image(self, image, return_image=False):
processed_image = [
processor(image)
for processor in self.processors
]
if return_image:
return processed_image
def process_image(self, image, processor_id=None):
if processor_id is None:
processed_image = [processor(image) for processor in self.processors]
else:
processed_image = [self.processors[processor_id](image)]
processed_image = torch.concat([
torch.Tensor(np.array(image_, dtype=np.float32) / 255).permute(2, 0, 1).unsqueeze(0)
for image_ in processed_image

View File

@@ -16,15 +16,15 @@ class Annotator:
if processor_id == "canny":
self.processor = CannyDetector()
elif processor_id == "depth":
self.processor = MidasDetector.from_pretrained(model_path)
self.processor = MidasDetector.from_pretrained(model_path).to("cuda")
elif processor_id == "softedge":
self.processor = HEDdetector.from_pretrained(model_path)
self.processor = HEDdetector.from_pretrained(model_path).to("cuda")
elif processor_id == "lineart":
self.processor = LineartDetector.from_pretrained(model_path)
self.processor = LineartDetector.from_pretrained(model_path).to("cuda")
elif processor_id == "lineart_anime":
self.processor = LineartAnimeDetector.from_pretrained(model_path)
self.processor = LineartAnimeDetector.from_pretrained(model_path).to("cuda")
elif processor_id == "openpose":
self.processor = OpenposeDetector.from_pretrained(model_path)
self.processor = OpenposeDetector.from_pretrained(model_path).to("cuda")
elif processor_id == "tile":
self.processor = None
else:

View File

@@ -7,7 +7,7 @@ import cupy as cp
class FastBlendSmoother:
def __init__(self):
self.batch_size = 8
self.window_size = 32
self.window_size = 64
self.ebsynth_config = {
"minimum_patch_size": 5,
"threads_per_block": 8,

View File

@@ -189,6 +189,7 @@ class ModelManager:
model.to(device)
else:
self.model[component].to(device)
torch.cuda.empty_cache()
def get_model_with_model_path(self, model_path):
for component in self.model_path:

View File

@@ -51,6 +51,9 @@ def lets_dance_with_long_video(
hidden_states = hidden_states * (num / (num + bias)) + hidden_states_updated * (bias / (num + bias))
hidden_states_output[i] = (hidden_states, num + 1)
if batch_id_ == num_frames:
break
# output
hidden_states = torch.stack([h for h, _ in hidden_states_output])
return hidden_states
@@ -195,10 +198,21 @@ class SDVideoPipeline(torch.nn.Module):
# Prepare ControlNets
if controlnet_frames is not None:
controlnet_frames = torch.stack([
self.controlnet.process_image(controlnet_frame).to(self.torch_dtype)
for controlnet_frame in progress_bar_cmd(controlnet_frames)
], dim=1)
if isinstance(controlnet_frames[0], list):
controlnet_frames_ = []
for processor_id in range(len(controlnet_frames)):
controlnet_frames_.append(
torch.stack([
self.controlnet.process_image(controlnet_frame, processor_id=processor_id).to(self.torch_dtype)
for controlnet_frame in progress_bar_cmd(controlnet_frames[processor_id])
], dim=1)
)
controlnet_frames = torch.concat(controlnet_frames_, dim=0)
else:
controlnet_frames = torch.stack([
self.controlnet.process_image(controlnet_frame).to(self.torch_dtype)
for controlnet_frame in progress_bar_cmd(controlnet_frames)
], dim=1)
# Denoise
for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)):