Create main.py

main.py

@@ -0,0 +1,252 @@
+import argparse
+import glob
+import os
+from pathlib import Path
+import uuid
+from src.pipelines.pipeline_animatediff_pix2pix import StableDiffusionInstructPix2PixPipeline
+from diffusers import EulerAncestralDiscreteScheduler
+import torch
+from src.models.unet import UNet3DConditionModel
+import numpy as np
+from PIL import Image
+import imageio
+
+def convert_frames_to_mp4(frames, filename, fps=30):
+    """Converts a list of PIL Image frames to an MP4 file.
+
+    Args:
+        frames: A list of PIL Image frames.
+        filename: The name of the MP4 file to save.
+        fps: Frames per second for the video.
+
+    Returns:
+        None
+    """
+    # Convert PIL Images to numpy arrays
+    numpy_frames = [np.array(frame) for frame in frames]
+    # Write frames to mp4
+    imageio.mimwrite(filename, numpy_frames, fps=fps)
+
+def convert_frames_to_gif(frames, filename, duration=100):
+    """Converts a list of PIL Image frames to a GIF file.
+
+    Args:
+        frames: A list of PIL Image frames.
+        filename: The name of the GIF file to save.
+        duration: Duration of each frame in milliseconds.
+
+    Returns:
+        None
+    """
+    frames[0].save(
+        filename, 
+        save_all=True, 
+        append_images=frames[1:], 
+        loop=0, 
+        duration=duration
+    )
+
+
+def convert_frames_to_gif_with_fps(frames, filename, fps=30):
+    """Converts a list of PIL Image frames to a GIF file using fps.
+
+    Args:
+        frames: A list of PIL Image frames.
+        filename: The name of the GIF file to save.
+        fps: Frames per second for the gif.
+
+    Returns:
+        None
+    """
+    duration = 1000 // fps
+    frames[0].save(
+        filename, 
+        save_all=True, 
+        append_images=frames[1:], 
+        loop=0, 
+        duration=duration
+    )
+
+
+def run(t2i_model,
+        prompt="", 
+        negative_prompt="", 
+        frame_count=16,
+        num_inference_steps=20,
+        guidance_scale=7.5,
+        image_guidance_scale=1.5,
+        width=512,
+        height=512,
+        dtype="float16",
+        output_frames_directory="output_frames",
+        output_video_directory="output_video",
+        output_gif_directory="output_gif",
+        motion_module="viddle/viddle-pix2pix-animatediff-v1.ckpt", 
+        init_image=None,
+        init_folder=None, 
+        seed=42,
+        fps=15,
+        no_save_frames=False,
+        no_save_video=False,
+        no_save_gif=False,
+        ):
+  scheduler_kwargs = {
+    "num_train_timesteps": 1000,
+    "beta_start": 0.00085,
+    "beta_end": 0.012,
+    "beta_schedule": "linear",
+  }
+
+  device = "cuda" if torch.cuda.is_available() else "cpu"
+  if dtype == "float16":
+     dtype = torch.float16
+     variant = "fp16"
+  elif dtype == "float32":
+      dtype = torch.float32
+      variant = "fp32"
+
+  unet_additional_kwargs = {
+    "in_channels": 8,
+    "unet_use_cross_frame_attention": False,
+    "unet_use_temporal_attention": False,
+    "use_motion_module": True,
+    "motion_module_resolutions": [1, 2, 4, 8],
+    "motion_module_mid_block": False,
+    "motion_module_decoder_only": False,
+    "motion_module_type": "Vanilla",
+    "motion_module_kwargs": {
+        "num_attention_heads": 8,
+        "num_transformer_block": 1,
+        "attention_block_types": ["Temporal_Self", "Temporal_Self"],
+        "temporal_position_encoding": True,
+        "temporal_position_encoding_max_len": 32,
+        "temporal_attention_dim_div": 1,
+    },
+  }
+
+  pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+      t2i_model,
+      scheduler=EulerAncestralDiscreteScheduler(**scheduler_kwargs),
+      safety_checker=None,
+      feature_extractor=None,
+      requires_safety_checker=False,
+      torch_dtype=dtype,
+      variant=variant,
+  ).to(device)
+
+  pipeline.unet = UNet3DConditionModel.from_pretrained_unet(pipeline.unet,
+                                                            unet_additional_kwargs=unet_additional_kwargs,
+                                                            ).to(device=device, dtype=dtype)
+  
+  pipeline.enable_vae_slicing()
+
+  motion_module_state_dict = torch.load(motion_module, map_location="cpu")
+  _, unexpected = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False)
+  assert len(unexpected) == 0
+
+  if init_image is not None and init_folder is None:
+    image = Image.open(init_image)
+    image = image.resize((width, height))
+  elif init_folder is not None and init_image is None:
+    image_paths = glob.glob(init_folder + "/*.png")
+    # add the jpgs
+    image_paths += glob.glob(init_folder + "/*.jpg")
+    image_paths.sort()
+    image_paths = image_paths[:frame_count]
+  
+    image = []
+
+    for image_path in image_paths:
+      image.append(Image.open(image_path).resize((width, height)))
+  else:
+    raise ValueError("Must provide either init_image or init_folder but not both")
+  
+  generator = torch.Generator(device=device).manual_seed(seed)
+
+  frames = pipeline(prompt=prompt, 
+                negative_prompt=negative_prompt, 
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                image_guidance_scale=image_guidance_scale,
+                image=image,
+                video_length=frame_count,
+                generator=generator,
+                )[0]
+
+  # create a uuid prefix for the output files
+  uuid_prefix = str(uuid.uuid4())
+
+  if not no_save_frames:
+    # Create output directory
+    Path(output_frames_directory).mkdir(parents=True, exist_ok=True)
+
+    # make the specific directory for this run
+    output_frames_directory = os.path.join(output_frames_directory, uuid_prefix)
+    Path(output_frames_directory).mkdir(parents=True, exist_ok=True)
+    # Save frames
+    for i, frame in enumerate(frames):
+      frame.save(os.path.join(output_frames_directory, f"{str(i).zfill(4)}.png"))
+
+  if not no_save_video:
+    # Create output directory
+    Path(output_video_directory).mkdir(parents=True, exist_ok=True)
+
+    convert_frames_to_mp4(frames, os.path.join(output_video_directory, f"{uuid_prefix}.mp4"), fps=fps)
+
+  if not no_save_gif:
+    # Create output directory
+    Path(output_gif_directory).mkdir(parents=True, exist_ok=True)
+
+    # Convert frames to GIF
+    convert_frames_to_gif(frames, os.path.join(output_gif_directory, f"{uuid_prefix}.gif"), duration=1000 // fps)
+
+
+if __name__ == "__main__":
+  argsparser = argparse.ArgumentParser()
+  argsparser.add_argument("--prompt", type=str, default="")
+  argsparser.add_argument("--negative_prompt", type=str, default="")
+  argsparser.add_argument("--frame_count", type=int, default=16)
+  argsparser.add_argument("--num_inference_steps", type=int, default=20)
+  argsparser.add_argument("--guidance_scale", type=float, default=7.5)
+  argsparser.add_argument("--image_guidance_scale", type=float, default=1.5)
+  argsparser.add_argument("--width", type=int, default=512)
+  argsparser.add_argument("--height", type=int, default=512)
+  argsparser.add_argument("--dtype", type=str, default="float16")
+  argsparser.add_argument("--output_frames_directory", type=str, default="output_frames")
+  argsparser.add_argument("--output_video_directory", type=str, default="output_videos")
+  argsparser.add_argument("--output_gif_directory", type=str, default="output_gifs")
+  argsparser.add_argument("--init_image", type=str, default=None)
+  argsparser.add_argument("--init_folder", type=str, default=None)
+  argsparser.add_argument("--motion_module", type=str, default="checkpoints/viddle-pix2pix-animatediff-v1.ckpt")
+  argsparser.add_argument("--t2i_model", type=str, default="timbrooks/instruct-pix2pix")
+  argsparser.add_argument("--seed", type=int, default=42)
+  argsparser.add_argument("--fps", type=int, default=15)
+  argsparser.add_argument("--no_save_frames", action="store_true", default=False)
+  argsparser.add_argument("--no_save_video", action="store_true", default=False)
+  argsparser.add_argument("--no_save_gif", action="store_true", default=False)
+  args = argsparser.parse_args()
+
+  run(t2i_model=args.t2i_model,
+      prompt=args.prompt, 
+      negative_prompt=args.negative_prompt, 
+      frame_count=args.frame_count,
+      num_inference_steps=args.num_inference_steps,
+      guidance_scale=args.guidance_scale,
+      width=args.width,
+      height=args.height,
+      dtype=args.dtype,
+      output_frames_directory=args.output_frames_directory,
+      output_video_directory=args.output_video_directory,
+      output_gif_directory=args.output_gif_directory,
+      motion_module=args.motion_module,
+      init_image=args.init_image,
+      init_folder=args.init_folder,
+      seed=args.seed,
+      fps=args.fps,
+      no_save_frames=args.no_save_frames,
+      no_save_video=args.no_save_video,
+      no_save_gif=args.no_save_gif,
+      )
+
+  
+