Update README.md

README.md

@@ -2,7 +2,7 @@
 license: apache-2.0
 ---
 
-This repository contains a pruned and partially reorganized version of [AniPortrait](https://github.com/Zejun-Yang/AniPortrait).
+This repository contains a pruned and partially reorganized version of [AniPortrait](https://github.com/Zejun-Yang/AniPortrait), with some new features.
 
 ```
 @misc{wei2024aniportrait,
@@ -15,6 +15,53 @@ This repository contains a pruned and partially reorganized version of [AniPortr
 }
 ```
 
+# Added Features
+
+## CPU Offloading
+
+In order to maximize VRAM, CPU offloading is made available and enabled by default when using the command-line.
+
+- In python, use `pipeline.enable_model_cpu_offload(gpu_id: int=0)` to enable it.
+- In the command-line, pass `--no-offload` or `-no` to disable it.
+
+## Video Reference
+
+In addition to using an image as reference, this repository permits using a video as a reference image using a modified temporally-aware reference U-Net.
+
+### Automatic Masked Composition
+
+When using a video reference, you can pass `paste_back=True` or `--paste-back/-pb` to automatically re-composite the generated face onto the reference video.
+
+### Automatic Face Cropping
+
+In addition to the above, you can pass `crop_to_face=True` or `--crop-to-face/-ctf` to automatically crop the input video to the detected face region prior to inference, then un-crop and compose the results. This means you can directly use this feature on HD video without needing to pre-process any of the inputs.
+
+### Examples
+
+The following examples were produced using the above-referenced procedure, then interpolated and post-processed using [enfugue](https://github.com/painebenjamin/app.enfugue.ai).
+
+They were made using the source videos at Full HD, using the following command-line format:
+
+```sh
+aniportrait video.mp4 --video video.mp4 --audio audio.mp3 --crop-to-face --num-inference-steps 50 -cfg 4.5
+```
+
+All videos sourced from [Pexels](https://www.pexels.com/), voices are AI-generated reading [Attention is All You Need (ARXIV:1706.03762)](https://arxiv.org/abs/1706.03762)
+
+<table>
+  <tr>
+    <td>
+      <video controls src="https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/nW0kiagVIP2Nr_K_SjJUP.mp4"</video>
+    </td>
+    <td>
+      <video controls src="https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/FOj4drfm2OLa-n7JxYyIj.mp4"></video>
+    </td>
+    <td>
+      <video controls src="https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/Tu2ac8Cuvaexo_S-kgpPv.mp4"></video>
+    </td>
+  </tr>
+</table>
+
 # Usage
 
 ## Installation
@@ -27,10 +74,10 @@ pip install git+https://github.com/painebenjamin/aniportrait.git
 
 ## Command-Line
 
-A command-line utilitiy `aniportrait` is installed with the package.
+A command-line utility `aniportrait` is installed with the package.
 
 ```sh
-Usage: aniportrait [OPTIONS] INPUT_IMAGE
+Usage: aniportrait [OPTIONS] INPUT_IMAGE_OR_VIDEO
 
   Run AniPortrait on an input image with a video, and/or audio file. When only
   a video file is provided, a video-to-video (face reenactment) animation is
@@ -63,6 +110,8 @@ Options:
                                   image height.
   -m, --model TEXT                HuggingFace model name.
   -nh, --no-half                  Do not use half precision.
+  -no, --no-offload               Do not offload to the CPU to preserve GPU
+                                  memory.
   -g, --gpu-id INTEGER            GPU ID to use.
   -sf, --model-single-file        Download and use a single file instead of a
                                   directory.
@@ -77,11 +126,36 @@ Options:
                                   aniportrait.safetensors]
   -rs, --remote-subfolder TEXT    Remote subfolder to download from when using
                                   the model-single-file option.
-  -c, --cache-dir DIRECTORY       Cache directory to download to. Default uses
+  -cd, --cache-dir DIRECTORY      Cache directory to download to. Default uses
                                   the huggingface cache.
   -o, --output FILE               Output file.  [default: output.mp4]
+  -pb, --paste-back               Paste the original background back in.
+  -pbcf, --paste-back-color-fix [adain|wavelet]
+                                  Color fix method to use when pasting back.
+                                  [default: wavelet]
+  -ctf, --crop-to-face            Crop the input to the face prior to
+                                  execution, then merge the cropped result
+                                  with the uncropped image. Implies --paste-
+                                  back.
+  -pop, --pose-output FILE        When passed, save the pose image(s) to this
+                                  file.
+  -mop, --mask-output FILE        When passed, save the mask image(s) to this
+                                  file.
+  -cop, --combined-output FILE    When passed, save the combined image(s) to
+                                  this file.
+  -mb, --mask-blur INTEGER        Amount of blur to apply to the mask when
+                                  using cropping or pasting.  [default: 15]
+  -md, --mask-dilate INTEGER      Amount of dilation to apply to the mask when
+                                  using cropping or pasting.  [default: 31]
+  -ms, --mask-slow                Use a slower, more accurate mask generation
+                                  method.
+  -lss, --leading-seconds-silence FLOAT
+                                  Seconds of silence to add to the beginning
+                                  of the audio.  [default: 0.0]
+  -tss, --trailing-seconds-silence FLOAT
+                                  Seconds of silence to add to the end of the
+                                  audio.  [default: 0.0]
   --help                          Show this message and exit.
-
 ```
 
 ## Python
@@ -196,6 +270,9 @@ pipeline.audio2vid(
     context_batch_size: int=1,
     interpolation_factor: int=1,
     use_long_video: bool=True,
+    pose_filename: Optional[str]=None,                                                                                                                                               
+    leading_seconds_silence: float=0.0,                                                                                                                                              
+    trailing_seconds_silence: float=0.0,     
     **kwargs: Any
 ) -> Pose2VideoPipelineOutput
 ```
@@ -205,7 +282,63 @@ Using an audio file, draw `fps` face pose images per second for the duration of
 - Optionally pass width/height to modify the size. Defaults to reference image size.
 - Optionally pass `use_long_video = false` to disable using the long video pipeline.
 - Optionally pass `reference_pose_image` to designate the pose of `reference_image`. When not passed, the pose of `reference_image` is automatically detected.
-- Optionally pass `video_length` to use this many frames. Default is the same as the length of the pose reference images.
+- Optionally pass `video_length` to use this many frames. Default is the same as the length of the pose reference images or the length of the audio frames (when translated to the proper FPS) - whichever is shorter.
+- Optionally pass `leading_seconds_silence` and/or `trailing_seconds_silence` to add silent frame(s) to the beginning and/or end of the audio. This will be adjusted for your passed or detected frame rate.
+- Optionally pass `pose_filename`, `mask_filename` and/or `combined_filename` to save the pose, mask and/or combined frames to this video file for debugging.
+- 
+### audiovid2vid
+
+```py
+pipeline.audiovid2vid(                                                                                                                                                                                                                                                                                                                                       
+    audio: str,                                                                                                                                                                      
+    reference_image: List[Image.Image],                                                                                                                                              
+    num_inference_steps: int=25,                                                                                                                                                     
+    guidance_scale: float=3.5,                                                                                                                                                       
+    fps: int=30,                                                                                                                                                                     
+    eta: float=0.0,                                                                                                                                                                  
+    reference_pose_image: Optional[Image.Image]=None,                                                                                                                                
+    pose_reference_images: Optional[List[Image.Image]]=None,                                                                                                                         
+    generation: Optional[Union[torch.Generator, List[torch.Generator]]]=None,                                                                                                        
+    output_type: Optional[str]="pil",                                                                                                                                                
+    return_dict: bool=True,                                                                                                                                                          
+    callback: Optional[Callable[[int, int, torch.FloatTensor], None]]=None,                                                                                                          
+    callback_steps: Optional[int]=None,                                                                                                                                              
+    context_schedule: str="uniform",                                                                                                                                                 
+    context_frames: int=16,                                                                                                                                                          
+    context_overlap: int=4,                                                                                                                                                          
+    context_batch_size: int=1,                                                                                                                                                       
+    interpolation_factor: int=1,                                                                                                                                                     
+    width: Optional[int]=None,                                                                                                                                                       
+    height: Optional[int]=None,                                                                                                                                                      
+    video_length: Optional[int]=None,                                                                                                                                                
+    use_long_video: bool=True,                                                                                                                                                       
+    paste_back: bool=True,                                                                                                                                                           
+    paste_back_color_fix: Optional[Literal["wavelet", "adain"]]="wavelet",                                                                                                           
+    crop_to_face: bool=False,                                                                                                                                                        
+    crop_to_face_target_size: Optional[int]=512,                                                                                                                                     
+    crop_to_face_padding: Optional[int]=64,                                                                                                                                          
+    mask_filename: Optional[str]=None,                                                                                                                                               
+    pose_filename: Optional[str]=None,                                                                                                                                               
+    combined_filename: Optional[str]=None,                                                                                                                                           
+    mask_dilate: Optional[int]=31,                                                                                                                                                   
+    mask_gaussian_kernel_size: Optional[int]=15,                                                                                                                                     
+    mask_first_frame: bool=True,                                                                                                                                                     
+    leading_seconds_silence: float=0.0,                                                                                                                                              
+    trailing_seconds_silence: float=0.0,                                                                                                                                             
+    **kwargs: Any                                                                                                                                                                    
+) -> Pose2VideoPipelineOutput:
+```
+
+Using an audio file, draw `fps` face pose images per second for the duration of the audio. Then, using those face pose images, render a video using a video as a reference.
+- Optionally pass width/height to modify the size. Defaults to reference image size.
+- Optionally pass `use_long_video = false` to disable using the long video pipeline.
+- Optionally pass `video_length` to use this many frames. Default is the same as the length of the pose reference images or the length of the audio frames (when translated to the proper FPS) - whichever is shorter.
+- Optionally pass `paste_back = true` to re-composite the output onto the input.
+- When using `paste_back`, the face is color-fixed when re-pasting in order to reduce visible difference. The default method is `wavelet`, pass `adain` or `None` for other options.
+- Optionally pass `crop_to_face = true` to crop all images to the face region (with padding) prior to diffusion. This implies `paste_back = true`.
+- When using `crop_to_face`, we must first identify where the faces are in a potentially large image - to do this we perform a slow tiled face detection across the whole image. In order to reduce time to generate additional masks, the default behavior is to only do this once, then use the faces from the first frame to know where to detect in subsequent frames. Set this to `false` to perform tiled faced detection on every frame of the input - this is slower but will allow for more variability in the frame.
+- Optionally pass `leading_seconds_silence` and/or `trailing_seconds_silence` to add silent frame(s) to the beginning and/or end of the audio. This will be adjusted for your passed or detected frame rate.
+- Optionally pass `pose_filename`, `mask_filename` and/or `combined_filename` to save the pose, mask and/or combined frames to this video file for debugging.
 
 ## Internals/Helpers