temporalvideo usage changes

Hi,

This is an updated version of the temporalvideo.py file that does the following things:

* Add argparser for easier basic usage
* Add dynamic lookup of controlnet models in case the names or commit ids of available models are different
* Ensure all image inputs are scaled to the same resolution - the optical flow maps were hard coded to 512x512 and would produce images that didn't line up with non square input images.
* Write out controlnet images to help debug resolution related errors. It's also interesting to see the output of the preprocessors for hed and openpose.

Hopefully you find this helpful. Thanks so much for your work on this topic, it's pretty awesome!

Jase

temporalvideo.py

@@ -4,6 +4,7 @@ import requests
 import json
 import cv2
 import numpy as np
+import re
 import sys
 import torch
 from PIL import Image
@@ -20,6 +21,27 @@ import cv2
 from torchvision.io import write_jpeg
 import pickle
 
+import argparse
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('prompt')
+    parser.add_argument('--negative-prompt', dest='negative_prompt', default="")
+
+    parser.add_argument('--init-image', dest='init_image', default="./init.png")
+    parser.add_argument('--input-dir', dest='input_dir', default="./Input_Images")
+    parser.add_argument('--output-dir', dest='output_dir', default="./output")
+
+    parser.add_argument('--width', default=512, type=int)
+    parser.add_argument('--height', default=512, type=int)
+
+    return parser.parse_args()
+
+
+args = get_args()
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -27,11 +49,9 @@ model = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).to(device
 model = model.eval()
 
 # Replace with the actual path to your image file and folder
-x_path = "./init.png"
-y_folder = "./Input_Images"
 
-output_folder = "output"
-os.makedirs(output_folder, exist_ok=True)
+os.makedirs(args.output_dir, exist_ok=True)
+
 
 def get_image_paths(folder):
     image_extensions = ("*.jpg", "*.jpeg", "*.png", "*.bmp")
@@ -40,7 +60,46 @@ def get_image_paths(folder):
         files.extend(glob.glob(os.path.join(folder, ext)))
     return sorted(files)
 
-y_paths = get_image_paths(y_folder)
+
+y_paths = get_image_paths(args.input_dir)
+
+
+def get_controlnet_models():
+    url = "http://localhost:7860/controlnet/model_list"
+
+    temporalnet_model = None
+    temporalnet_re = re.compile("^temporalnetversion2 \[.{8}\]")
+
+    hed_model = None
+    hed_re = re.compile("^control_.*hed.* \[.{8}\]")
+
+    openpose_model = None
+    openpose_re = re.compile("^control_.*openpose.* \[.{8}\]")
+
+    response = requests.get(url)
+    if response.status_code == 200:
+        models = json.loads(response.content)
+    else:
+        raise Exception("Unable to list models from the SD Web API! "
+                        "Is it running and is the controlnet extension installed?")
+
+    for model in models['model_list']:
+        if temporalnet_model is None and temporalnet_re.match(model):
+            temporalnet_model = model
+        elif hed_model is None and hed_re.match(model):
+            hed_model = model
+        elif openpose_model is None and openpose_re.match(model):
+            openpose_model = model
+
+    assert temporalnet_model is not None, "Unable to find the temporalnet2 model!  Ensure it's copied into the stable-diffusion-webui/extensions/models directory!"
+    assert hed_model is not None, "Unable to find the hed_model model!  Ensure it's copied into the stable-diffusion-webui/extensions/models directory!"
+    assert openpose_model is not None, "Unable to find the openpose model!  Ensure it's copied into the stable-diffusion-webui/extensions/models directory!"
+
+    return temporalnet_model, hed_model, openpose_model
+
+
+TEMPORALNET_MODEL, HED_MODEL, OPENPOSE_MODEL = get_controlnet_models()
+
 
 def send_request(last_image_path, optical_flow_path,current_image_path):
     url = "http://localhost:7860/sdapi/v1/img2img"
@@ -51,7 +110,6 @@ def send_request(last_image_path, optical_flow_path,current_image_path):
     # Load and process the last image
     last_image = cv2.imread(last_image_path)
     last_image = cv2.cvtColor(last_image, cv2.COLOR_BGR2RGB)
-    last_image = cv2.resize(last_image, (512, 512))
 
     # Load and process the optical flow image
     flow_image = cv2.imread(optical_flow_path)
@@ -79,31 +137,39 @@ def send_request(last_image_path, optical_flow_path,current_image_path):
         "inpainting_mask_invert": 1,
         "resize_mode": 0,
         "denoising_strength": 0.4,
-        "prompt": "1girl, woman",
-        "negative_prompt": "",
+        "prompt": args.prompt,
+        "negative_prompt": args.negative_prompt,
         "alwayson_scripts": {
             "ControlNet":{
                 "args": [
                     {
                         "input_image": current_image,
                         "module": "hed",
-                        "model": "control_hed-fp16 [13fee50b]",
+                        "model": HED_MODEL,
                         "weight": 0.7,
                         "guidance": 1,
+                        "pixel_perfect": True,
+                        "resize_mode": 0,
                    },
                     {
                         "input_image": encoded_image,
-                        "model": "temporalnetversion2 [b146ac48]",
+                        "model": TEMPORALNET_MODEL,
                         "module": "none",
                         "weight": 0.6,
                         "guidance": 1,
+                        # "processor_res": 512,
+                        "threshold_a": 64,
+                        "threshold_b": 64,
+                        "resize_mode": 0,
                     },
                     {
                         "input_image": current_image,
-                        "model": "control_v11p_sd15_openpose [cab727d4]",
+                        "model": OPENPOSE_MODEL,
                         "module": "openpose_full",
                         "weight": 0.7,
-                        "guidance":1,
+                        "guidance": 1,
+                        "pixel_perfect": True,
+                        "resize_mode": 0,
                     }
                     
                   
@@ -118,8 +184,8 @@ def send_request(last_image_path, optical_flow_path,current_image_path):
         "n_iter": 1,
         "steps": 20,
         "cfg_scale": 6,
-        "width": 512,
-        "height": 512,
+        "width": args.width,
+        "height": args.height,
         "restore_faces": True,
         "include_init_images": True,
         "override_settings": {},
@@ -164,25 +230,18 @@ def infer(frameA, frameB):
         img2_batch = F.resize(img2_batch, size=[512, 512])
         return transforms(img1_batch, img2_batch)
 
-
     img1_batch, img2_batch = preprocess(img1_batch, img2_batch)
 
-
     list_of_flows = model(img1_batch.to(device), img2_batch.to(device))
 
-    predicted_flows = list_of_flows[-1]
-
+    predicted_flow = list_of_flows[-1][0]
+    opitcal_flow_path = os.path.join(args.output_dir, f"flow_{i}.png")
 
-    #flow_imgs = flow_to_image(predicted_flows)
+    flow_img = flow_to_image(predicted_flow).to("cpu")
+    flow_img = F.resize(flow_img, size=[args.height, args.width])
 
-    #print(flow_imgs)
+    write_jpeg(flow_img, opitcal_flow_path)
 
-    predicted_flow = list_of_flows[-1][0]
-    opitcal_flow_path = os.path.join(output_folder, f"flow_{i}.png")
-    flow_img = flow_to_image(predicted_flow).to("cpu")
-    write_jpeg(flow_img,opitcal_flow_path)
-    
-    
     return opitcal_flow_path
 
 output_images = []
@@ -190,13 +249,13 @@ output_paths = []
 
 # Initialize with the first image path
 
-result = x_path
-output_image_path = os.path.join(output_folder, f"output_image_0.png")
+result = args.init_image
+output_image_path = os.path.join(args.output_dir, f"output_image_0.png")
 
 #with open(output_image_path, "wb") as f:
    # f.write(result)
     
-last_image_path = x_path
+last_image_path = args.init_image
 for i in range(1, len(y_paths)):
     # Use the last image path and optical flow map to generate the next input
     optical_flow = infer(y_paths[i - 1], y_paths[i])
@@ -204,10 +263,14 @@ for i in range(1, len(y_paths)):
     # Modify your send_request to use the last_image_path
     result = send_request(last_image_path, optical_flow, y_paths[i])
     data = json.loads(result)
-    encoded_image = data["images"][0]
-    output_image_path = os.path.join(output_folder, f"output_image_{i}.png")
-    last_image_path = output_image_path
-    with open(output_image_path, "wb") as f:
-       f.write(base64.b64decode(encoded_image))
-    print(f"Written data for frame {i}:")
 
+    for j, encoded_image in enumerate(data["images"]):
+        if j == 0:
+            output_image_path = os.path.join(args.output_dir, f"output_image_{i}.png")
+            last_image_path = output_image_path
+        else:
+            output_image_path = os.path.join(args.output_dir, f"controlnet_image_{j}_{i}.png")
+
+        with open(output_image_path, "wb") as f:
+           f.write(base64.b64decode(encoded_image))
+    print(f"Written data for frame {i}:")