Fix pipeline

.gitignore

@@ -1,3 +1,5 @@
 env/
 __pycache__/
-assets/
+assets/
+*.png
+*.mp4

bapp.py

@@ -4,10 +4,12 @@ from pathlib import Path
 import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
+from icecream import ic
 
 from visual_anagrams.views import get_views, VIEW_MAP_NAMES
 from visual_anagrams.samplers import sample_stage_1, sample_stage_2
 from visual_anagrams.utils import add_args, save_illusion, save_metadata
+from visual_anagrams.animate import animate_two_view
 
 stage_1 = DiffusionPipeline.from_pretrained(
                 "DeepFloyd/IF-I-M-v1.0",
@@ -31,23 +33,26 @@ def generate_content(
     num_inference_steps,
     seed
 ):
-    prompts = [prompt_for_original, prompt_for_transformed]
-    prompt_embeds = [stage_1.encode_prompt(f'{style} {p}'.strip()) for p in [prompts]]
+    prompts = [f'{style} {p}'.strip() for p in [prompt_for_original, prompt_for_transformed]]
+    prompt_embeds = [stage_1.encode_prompt(p) for p in prompts]
     prompt_embeds, negative_prompt_embeds = zip(*prompt_embeds)
     prompt_embeds = torch.cat(prompt_embeds)
     negative_prompt_embeds = torch.cat(negative_prompt_embeds)
 
-    views = ['identity', transformation]
+    views = ['identity', VIEW_MAP_NAMES[transformation]]
     views = get_views(views)
 
     generator = torch.manual_seed(seed)
+
+    print("Sample stage 1")
     image = sample_stage_1(stage_1,
-                            prompt_embeds,
-                            negative_prompt_embeds,
-                            views,
-                            num_inference_steps=num_inference_steps,
-                            generator=generator)
+                           prompt_embeds,
+                           negative_prompt_embeds,
+                           views,
+                           num_inference_steps=num_inference_steps,
+                           generator=generator)
 
+    print("Sample stage 2")
     image = sample_stage_2(stage_2,
                            image,
                            prompt_embeds,
@@ -55,8 +60,16 @@ def generate_content(
                            views,
                            num_inference_steps=num_inference_steps,
                            generator=generator)
+    save_illusion(image, views, Path(""))
     
-    return image, image_transformed, transformation_gif
+    size = image.shape[-1]
+    animate_two_view(
+        f"sample_{size}.png",
+        views[1],
+        prompts[0],
+        prompts[1],
+    )
+    return 'tmp.mp4', f"sample_{size}.png", f"sample_{size}.views.png"
 
 
 choices = list(VIEW_MAP_NAMES.keys())
@@ -64,13 +77,13 @@ gradio_app = gr.Interface(
     fn=generate_content,
     inputs=[
         gr.Textbox(label="Style", placeholder="an oil painting of"),
-        gr.Textbox(label="Prompt for original view", placeholder="a penguin"),
-        gr.Textbox(label="Prompt for transformed view", placeholder="a giraffe"),
+        gr.Textbox(label="Prompt for original view", placeholder="a dress"),
+        gr.Textbox(label="Prompt for transformed view", placeholder="an old man"),
         gr.Dropdown(label="View transformation", choices=choices, value=choices[0]),
-        gr.Number(label="Number of diffusion steps", value=30, step=1, minimum=1, maximum=100),
+        gr.Number(label="Number of diffusion steps", value=100, step=1, minimum=1, maximum=300),
         gr.Number(label="Random seed", value=0, step=1, minimum=0, maximum=100000)
     ],
-    outputs=[gr.Image(label="Illusion"), gr.Image(label="Original"), gr.Image(label="Transformed")],
+    outputs=[gr.Video(label="Illusion"), gr.Image(label="Original"), gr.Image(label="Transformed")],
 )
 
 

requirements.txt

@@ -2,9 +2,15 @@ accelerate
 diffusers
 einops
 gradio
+icecream
+imageio
+imageio[ffmpeg]
+imageio[pyav]
+opencv-python
 safetensors
 sentencepiece
 transformers
 torch
 torchvision
+tqdm
 xformers

test.py

@@ -0,0 +1,12 @@
+from bapp import generate_content
+
+if __name__ == "__main__":
+    print(generate_content(
+        "a painting of",
+        "vases",
+        "a sloth",
+        "Flip",
+        1,
+        0
+    ))
+    

test_video.py

@@ -0,0 +1,11 @@
+from visual_anagrams.animate import animate_two_view
+from visual_anagrams.views import get_views
+
+if __name__ == "__main__":
+    animate_two_view(
+        "sample_256.png",
+        get_views(["identity", "flip"])[1],
+        "a painting of vases",
+        "a painting of a sloth",
+        save_video_path="tmp3.mp4",
+    )

animate.py → visual_anagrams/animate.py

@@ -1,3 +1,4 @@
+import cv2
 from tqdm import tqdm
 import numpy as np
 from PIL import Image, ImageDraw, ImageFont
@@ -13,12 +14,11 @@ def draw_text(image, text, fill=(0,0,0), frame_size=384, im_size=256):
     image = image.copy()
 
     # Font info
-    font_path = get_courier_font_path()
     font_size = 16
 
     # Make PIL objects
     draw = ImageDraw.Draw(image)
-    font = ImageFont.truetype(font_path, font_size)
+    font = ImageFont.load_default()
     
     # Center text horizontally, and vertically between
     # illusion bottom and frame bottom
@@ -117,10 +117,19 @@ def animate_two_view(
 
     # Convert PIL images to numpy arrays
     image_array = [imageio.core.asarray(frame) for frame in frames]
-
-    # Save as video
-    print('Making video...')
-    imageio.mimsave(save_video_path, image_array, fps=30)
+    f = image_array[0]
+    print(f.dtype)
+    print(f.shape)
+    print(frame_size)
+    print(np.min(f), np.max(f))
+    print(len(image_array))
+
+    # Save as video using opencv
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video = cv2.VideoWriter(save_video_path, fourcc, 30, (frame_size, frame_size))
+    for frame in image_array:
+        video.write(frame)
+    video.release()
 
 
 

visual_anagrams/samplers.py

@@ -1,4 +1,5 @@
 from tqdm import tqdm
+from icecream import ic
 
 import torch
 import torch.nn.functional as F
@@ -42,8 +43,9 @@ def sample_stage_1(model,
         device,
         generator,
     )
+    # ic(noisy_images.shape)
 
-    for i, t in enumerate(tqdm(timesteps)):
+    for i, t in tqdm(enumerate(timesteps)):
         # Apply views to noisy_image
         viewed_noisy_images = []
         for view_fn in views:
@@ -56,6 +58,7 @@ def sample_stage_1(model,
         model_input = model.scheduler.scale_model_input(model_input, t)
 
         # Predict noise estimate
+        # print("Predicting noise estimate")
         noise_pred = model.unet(
             model_input,
             t,
@@ -63,9 +66,11 @@ def sample_stage_1(model,
             cross_attention_kwargs=None,
             return_dict=False,
         )[0]
+        # ic(noise_pred.shape)
 
         # Extract uncond (neg) and cond noise estimates
         noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        # ic(noise_pred_uncond.shape)
 
         # Invert the unconditional (negative) estimates
         inverted_preds = []
@@ -73,6 +78,7 @@ def sample_stage_1(model,
             inverted_pred = view.inverse_view(pred)
             inverted_preds.append(inverted_pred)
         noise_pred_uncond = torch.stack(inverted_preds)
+        # ic(noise_pred_uncond.shape)
 
                                 # Invert the conditional estimates
         inverted_preds = []
@@ -80,11 +86,13 @@ def sample_stage_1(model,
             inverted_pred = view.inverse_view(pred)
             inverted_preds.append(inverted_pred)
         noise_pred_text = torch.stack(inverted_preds)
+        # ic(noise_pred_text.shape)
 
         # Split into noise estimate and variance estimates
         noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1)
         noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1)
         noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        # ic(noise_pred.shape)
 
         # Reduce predicted noise and variances
         noise_pred = noise_pred.view(-1,num_prompts,3,64,64)
@@ -98,11 +106,14 @@ def sample_stage_1(model,
         else:
             raise ValueError('Reduction must be either `mean` or `alternate`')
         noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)
+        # ic(noise_pred.shape)
 
+        # ic(t.shape)
         # compute the previous noisy sample x_t -> x_t-1
         noisy_images = model.scheduler.step(
-            noise_pred, t, noisy_images, generator=generator, return_dict=False
+            noise_pred.to('cuda'), t, noisy_images.to('cuda'), generator=generator, return_dict=False
         )[0]
+        # ic(noisy_images.shape)
 
     # Return denoised images
     return noisy_images
@@ -149,34 +160,34 @@ def sample_stage_2(model,
         prompt_embeds.dtype,
         device,
         generator,
-    )
+    ).to('cuda')
 
     # Prepare upscaled image and noise level
     image = model.preprocess_image(image, num_images_per_prompt, device)
-    upscaled = F.interpolate(image, (height, width), mode="bilinear", align_corners=True)
+    upscaled = F.interpolate(image.to('cuda'), (height, width), mode="bilinear", align_corners=True).to('cuda')
 
     noise_level = torch.tensor([noise_level] * upscaled.shape[0], device=upscaled.device)
     noise = randn_tensor(upscaled.shape, generator=generator, device=upscaled.device, dtype=upscaled.dtype)
-    upscaled = model.image_noising_scheduler.add_noise(upscaled, noise, timesteps=noise_level)
+    upscaled = model.image_noising_scheduler.add_noise(upscaled, noise, timesteps=noise_level).to('cuda')
 
     # Condition on noise level, for each model input
-    noise_level = torch.cat([noise_level] * num_prompts * 2)
+    noise_level = torch.cat([noise_level] * num_prompts * 2).to('cuda')
 
     # Denoising Loop
     for i, t in enumerate(tqdm(timesteps)):
         # Cat noisy image with upscaled conditioning image
-        model_input = torch.cat([noisy_images, upscaled], dim=1)
+        model_input = torch.cat([noisy_images, upscaled], dim=1).to('cuda')
 
         # Apply views to noisy_image
         viewed_inputs = []
         for view_fn in views:
             viewed_inputs.append(view_fn.view(model_input[0]))
-        viewed_inputs = torch.stack(viewed_inputs)
+        viewed_inputs = torch.stack(viewed_inputs).to('cuda')
 
         # Duplicate inputs for CFG
         # Model input is: [ neg_0, neg_1, ..., pos_0, pos_1, ... ]
-        model_input = torch.cat([viewed_inputs] * 2)
-        model_input = model.scheduler.scale_model_input(model_input, t)
+        model_input = torch.cat([viewed_inputs] * 2).to('cuda')
+        model_input = model.scheduler.scale_model_input(model_input, t).to('cuda')
 
         # predict the noise residual
         noise_pred = model.unet(