Spaces:
Runtime error
Runtime error
File size: 2,064 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
from scripts.external_code import ControlNetUnit
from scripts.logging import logger
from modules.processing import StableDiffusionProcessing
from modules import shared
def add_animate_diff_batch_input(
p: StableDiffusionProcessing, unit: ControlNetUnit
) -> ControlNetUnit:
"""AnimateDiff + ControlNet batch processing."""
assert unit.is_animate_diff_batch
batch_parameters = unit.batch_images.split("\n")
batch_image_dir = batch_parameters[0]
logger.info(
f"AnimateDiff + ControlNet {unit.module} receive the following parameters:"
)
logger.info(f"\tbatch control images: {batch_image_dir}")
for ad_cn_batch_parameter in batch_parameters[1:]:
if ad_cn_batch_parameter.startswith("mask:"):
unit.batch_mask_dir = ad_cn_batch_parameter[len("mask:") :].strip()
logger.info(f"\tbatch control mask: {unit.batch_mask_dir}")
elif ad_cn_batch_parameter.startswith("keyframe:"):
unit.batch_keyframe_idx = ad_cn_batch_parameter[len("keyframe:") :].strip()
unit.batch_keyframe_idx = [
int(b_i.strip()) for b_i in unit.batch_keyframe_idx.split(",")
]
logger.info(f"\tbatch control keyframe index: {unit.batch_keyframe_idx}")
batch_image_files = shared.listfiles(batch_image_dir)
for batch_modifier in getattr(unit, "batch_modifiers", []):
batch_image_files = batch_modifier(batch_image_files, p)
unit.batch_image_files = batch_image_files
unit.image = []
for idx, image_path in enumerate(batch_image_files):
mask_path = None
if getattr(unit, "batch_mask_dir", None) is not None:
batch_mask_files = shared.listfiles(unit.batch_mask_dir)
if len(batch_mask_files) >= len(batch_image_files):
mask_path = batch_mask_files[idx]
else:
mask_path = batch_mask_files[0]
unit.image.append(
{
"image": image_path,
"mask": mask_path,
}
)
return unit
|