simplification to flow generation only

app.py

@@ -17,9 +17,6 @@ using our implementation of the RAFT model. We will also see how to convert the
 predicted flows to RGB images for visualization.
 """
 
-from diffusers import DiffusionPipeline, ControlNetModel
-from diffusers import UniPCMultistepScheduler
-
 import cv2
 import numpy as np
 import os
@@ -42,53 +39,6 @@ from scipy.interpolate import LinearNDInterpolator
 from imageio import imread, imwrite
 
 
-# Constants
-low_threshold = 100
-high_threshold = 200
-
-# Models
-controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-pipe = DiffusionPipeline.from_pretrained(
-    "mikesmodels/Waltz_with_Bashir_Diffusion", controlnet=controlnet, custom_pipeline="stable_diffusion_controlnet_img2img", safety_checker=None, torch_dtype=torch.float16
-)
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-
-# This command loads the individual model components on GPU on-demand. So, we don't
-# need to explicitly call pipe.to("cuda").
-pipe.enable_model_cpu_offload()
-
-pipe.enable_xformers_memory_efficient_attention()
-
-# Generator seed,
-generator = torch.Generator('cuda').manual_seed(int(123456))
-
-def get_canny_filter(image):
-    if not isinstance(image, np.ndarray):
-        image = np.array(image) 
-        
-    image = cv2.Canny(image, low_threshold, high_threshold)
-    image = image[:, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    canny_image = Image.fromarray(image)
-    return canny_image
-
-
-def generate_images(prompt, canny_image, image):
-    
-    output = pipe(
-        controlnet_conditioning_image=canny_image,
-        image = image,
-        prompt = prompt,
-        generator=generator,
-        num_images_per_prompt=1,
-        num_inference_steps=20,
-    )
-    all_outputs = []
-    all_outputs.append(canny_image)
-    for image in output.images:
-        all_outputs.append(image)
-    return all_outputs
-
 
 def write_flo(flow, filename):
     """
@@ -125,20 +75,7 @@ def infer():
     #frames, _, _ = read_video(str("./spacex.mp4"), output_format="TCHW")
     #print(f"FRAME BEFORE stack: {frames[100]}")
     
-    prompt = "wltzwthbshr basketball player"
-    
-    pil2diff_img = Image.open("./basket1.jpg")
-    canny_image = get_canny_filter(pil2diff_img)
-    diffused_img = generate_images(prompt, canny_image, pil2diff_img)
-    print(f"DIFFUSED IMG: {diffused_img[1]}")
-    
-    diffused_img[1].save("diffused_input1.jpg")
     
-    pil2diff_img2 = Image.open("./frame2.jpg")
-    canny_image2 = get_canny_filter(pil2diff_img2)
-
-    canny_image.save("canny1.jpg")
-    canny_image2.save("canny2.jpg")
     input_frame_1 = read_image(str("./basket1.jpg"), ImageReadMode.UNCHANGED)
     print(f"FRAME 1: {input_frame_1}")
     input_frame_2 = read_image(str("./basket2.jpg"), ImageReadMode.UNCHANGED)
@@ -233,94 +170,8 @@ def infer():
     write_jpeg(flow_img, f"predicted_flow.jpg")
     
     flo_file = write_flo(predicted_flow, "flofile.flo")
-      
-    # define a transform to convert a tensor to PIL image
-    transform = T.ToPILImage()
     
-    # convert the tensor to PIL image using above transform
-    #img = transform(frames[1])
-    img = transform(input_frame_2)
-    img = img.resize((960, 520))
-    # display the PIL image
-    #img.show()
-    frame2pil = np.array(img.convert('RGB'))
-    print(f"frame1pil: {frame2pil}")
-    print(f"frame1pil shape: {frame2pil.shape}")
-    print(f"frame1pil dtype: {frame2pil.dtype}")
-    img.save('raw_frame2.jpg')
-
-    # convert the tensor diffused to PIL image using above transform
-    #img = transform(frames[1])
-    img_diff = transform(input_diffused)
-    img_diff = img_diff.resize((960, 520))
-    # display the PIL image
-    #img.show()
-    diffpil = np.array(img_diff.convert('RGB'))
-    print(f"frame1pil: {diffpil}")
-    print(f"frame1pil shape: {diffpil.shape}")
-    print(f"frame1pil dtype: {diffpil.dtype}")
-    img_diff.save('diffused_resized.jpg')
-    
-
-    numpy_array_flow = predicted_flow.permute(1, 2, 0).detach().cpu().numpy()
-    print(f"numpy_array_flow: {numpy_array_flow}")
-    print(f"numpy_array_flow shape: {numpy_array_flow.shape}")
-    print(f"numpy_array_flow dtype: {numpy_array_flow.dtype}")
-
-    h, w = numpy_array_flow.shape[:2]
-    numpy_array_flow = numpy_array_flow.copy()
-    numpy_array_flow[:, :, 0] += np.arange(w)
-    numpy_array_flow[:, :, 1] += np.arange(h)[:, np.newaxis]
-    # print('flow stats', flow.max(), flow.min(), flow.mean())
-    # print(flow)
-    numpy_array_flow*=1.
-    # print('flow stats mul', flow.max(), flow.min(), flow.mean())
-    # res = cv2.remap(img, flow, None, cv2.INTER_LINEAR)
-    res = cv2.remap(diffpil, numpy_array_flow, None, cv2.INTER_LANCZOS4)
-    print(res)
-    
-    res = Image.fromarray(res)
-    res.save('warped.jpg')
-
-    blend2 = Image.open('raw_frame2.jpg')
-    blend2 = Image.blend(res,blend2,0.5)
-    blend2.save("blended2.jpg")
-
-    pil2diff_blend = Image.open("warped.jpg")
-    #pil2diff_blend = Image.open("./basket2.jpg")
-    canny_image = get_canny_filter(pil2diff_blend)
-    diffused_blend = generate_images(prompt, canny_image, pil2diff_blend)
-    print(f"DIFFUSED IMG: {diffused_blend[1]}")
-    
-    diffused_blend[1].save("diffused_blended_2.jpg")
-    
-    return "done", "predicted_flow.jpg", ["flofile.flo"], "diffused_input1.jpg", "diffused_blended_2.jpg", 'warped.jpg', "blended2.jpg"
-####################################
-# Bonus: Creating GIFs of predicted flows
-# ---------------------------------------
-# In the example above we have only shown the predicted flows of 2 pairs of
-# frames. A fun way to apply the Optical Flow models is to run the model on an
-# entire video, and create a new video from all the predicted flows. Below is a
-# snippet that can get you started with this. We comment out the code, because
-# this example is being rendered on a machine without a GPU, and it would take
-# too long to run it.
-
-# from torchvision.io import write_jpeg
-# for i, (img1, img2) in enumerate(zip(frames, frames[1:])):
-#     # Note: it would be faster to predict batches of flows instead of individual flows
-#     img1, img2 = preprocess(img1, img2)
-
-#     list_of_flows = model(img1.to(device), img2.to(device))
-#     predicted_flow = list_of_flows[-1][0]
-#     flow_img = flow_to_image(predicted_flow).to("cpu")
-#     output_folder = "/tmp/"  # Update this to the folder of your choice
-#     write_jpeg(flow_img, output_folder + f"predicted_flow_{i}.jpg")
-
-####################################
-# Once the .jpg flow images are saved, you can convert them into a video or a
-# GIF using ffmpeg with e.g.:
-#
-# ffmpeg -f image2 -framerate 30 -i predicted_flow_%d.jpg -loop -1 flow.gif
+    return "done", "predicted_flow.jpg", ["flofile.flo"]
 
       
-gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(label="flow"), gr.Files(), gr.Image(label="diffused input"), gr.Image(), gr.Image(label="warped flow to img2"), gr.Image(label="blended result to diffuse")]).launch()
+gr.Interface(fn=infer, inputs=[], outputs=[gr.Textbox(), gr.Image(label="flow"), gr.Files()]).launch()