- direction_symbol = colors[direction_symbol].permute(0, 3, 1, 2)
- separator = torch.full((direction.size(0), 3, height * upscale - 1, 1), 0)
-
- for n in range(direction_symbol.size(0)):
- if direction[n] == token_forward:
- for k in range(upscale):
- direction_symbol[
- n,
- :,
- (height * upscale) // 2 - upscale // 2 + k,
- 3 + upscale // 2 - abs(k - upscale // 2),
- ] = 0
- elif direction[n] == token_backward:
- for k in range(upscale):
- direction_symbol[
- n,
- :,
- (height * upscale) // 2 - upscale // 2 + k,
- 3 + abs(k - upscale // 2),
- ] = 0
- else:
- for k in range(2, upscale - 2):
- direction_symbol[
- n, :, (height * upscale) // 2 - upscale // 2 + k, k
- ] = 0
- direction_symbol[
- n,
- :,
- (height * upscale) // 2 - upscale // 2 + k,
- upscale - 1 - k,
- ] = 0
-
- return torch.cat(
- [
- frame2img(f_first, height, width, upscale),
- separator,
- direction_symbol,
- separator,
- frame2img(f_second, height, width, upscale),
- ],
- dim=3,
+
+ separator[:, :, 0] = 0
+ separator[:, :, h - 1] = 0
+
+ for k in range(1, 2 * marker_size - 8):
+ i = k - (marker_size - 4)
+ j = marker_size - 5 - abs(i)
+ separator[:, :, h // 2 - 1 + i, 2 + j] = 0
+ separator[:, :, h // 2 - 1 + i + 1, 2 + j] = 0
+
+ img = torch.cat([img_prompts, separator, img_answers], dim=3)
+
+ image_name = os.path.join(result_dir, filename)
+ torchvision.utils.save_image(
+ img.float() / 255.0, image_name, nrow=6, padding=margin * 4, pad_value=1.0