Release Omni-Zero

README.md

@@ -1,6 +1,7 @@
 # Omni-Zero: A diffusion pipeline for zero-shot stylized portrait creation. 
 - [x] Release single person code
 - [ ] Release couples code
+- [ ] Add LoRA support
 
 ## Use Omni-Zero in [fal.ai](https://fal.ai) Workflows [https://fal.ai/dashboard/workflows/okaris/omni-zero](https://fal.ai/dashboard/workflows/okaris/omni-zero)
 ![Omni-Zero](https://github.com/okaris/omni-zero/assets/1448702/2ccbdf24-eb41-4a85-975e-af701fc4a879)
@@ -26,4 +27,4 @@ python demo.py
 - Special thanks to [fal.ai](https://fal.ai) for providing compute for the research and hosting
 - This project wouldn't be possible without the great work of the [InstantX Team](https://github.com/InstantID)
 - Thanks to [@fofrAI](http://twitter.com/fofrAI) for inspiring me with his [face-to-many workflow](https://github.com/fofr/cog-face-to-many)
-- Thanks to Matteo ([@cubiq](https://twitter.com/cubiq])) for creating the ComfyUI nodes for IP-Adapter
+- Thanks to Matteo ([@cubiq](https://twitter.com/cubiq])) for creating the ComfyUI nodes for IP-Adapter

omni_zero.py

@@ -57,6 +57,7 @@ class OmniZeroSingle():
         self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++", final_sigmas_type="zero")
 
         self.pipeline.load_ip_adapter(["okaris/ip-adapter-instantid", "h94/IP-Adapter", "h94/IP-Adapter"], subfolder=[None, "sdxl_models", "sdxl_models"], weight_name=["ip-adapter-instantid.bin", "ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus_sdxl_vit-h.safetensors"])
+   
     def get_largest_face_embedding_and_kps(self, image, target_image=None):
         face_info = self.face_analysis.get(cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR))
         if len(face_info) == 0:
@@ -156,4 +157,66 @@ class OmniZeroSingle():
             seed=seed,
         ).images
 
-        return images
+        return images
+    
+class OmniZeroCouple():
+    def __init__(self,
+        base_model="stabilityai/stable-diffusion-xl-base-1.0",
+    ):
+        snapshot_download("okaris/antelopev2", local_dir="./models/antelopev2")
+        self.face_analysis = FaceAnalysis(name='antelopev2', root='./', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
+
+        dtype = torch.float16
+
+        ip_adapter_plus_image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+            "h94/IP-Adapter", 
+            subfolder="models/image_encoder",
+            torch_dtype=dtype,
+        ).to("cuda")
+
+        zoedepthnet_path = "okaris/zoe-depth-controlnet-xl"
+        zoedepthnet = ControlNetModel.from_pretrained(zoedepthnet_path,torch_dtype=dtype).to("cuda")
+
+        identitiynet_path = "okaris/face-controlnet-xl"
+        identitynet = ControlNetModel.from_pretrained(identitiynet_path, torch_dtype=dtype).to("cuda")
+
+        self.zoe_depth_detector = ZoeDetector.from_pretrained("lllyasviel/Annotators").to("cuda")
+
+        self.pipeline = OmniZeroPipeline.from_pretrained(
+            base_model,
+            controlnet=[identitynet, zoedepthnet],
+            torch_dtype=dtype,
+            image_encoder=ip_adapter_plus_image_encoder,
+        ).to("cuda")
+
+        config = self.pipeline.scheduler.config
+        config["timestep_spacing"] = "trailing"
+        self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(config, use_karras_sigmas=True, algorithm_type="sde-dpmsolver++", final_sigmas_type="zero")
+
+        self.pipeline.load_ip_adapter(["okaris/ip-adapter-instantid", "okaris/ip-adapter-instantid", "h94/IP-Adapter", "h94/IP-Adapter"], subfolder=[None, None, "sdxl_models", "sdxl_models"], weight_name=["ip-adapter-instantid.bin", "ip-adapter-instantid.bin", "ip-adapter-plus_sdxl_vit-h.safetensors", "ip-adapter-plus_sdxl_vit-h.safetensors"])
+   
+    def generate(self,
+        seed=42,
+        prompt="A person",
+        negative_prompt="blurry, out of focus",
+        guidance_scale=3.0,
+        number_of_images=1,
+        number_of_steps=10,
+        base_image=None,
+        base_image_strength=0.15,
+        composition_image=None,
+        composition_image_strength=1.0,
+        style_image=None,
+        style_image_strength=1.0,
+        style_image_2=None,
+        style_image_strength_2=1.0,
+        identity_image=None,
+        identity_image_strength=1.0,
+        identity_image_2=None,
+        identity_image_strength_2=1.0,
+        depth_image=None,
+        depth_image_strength=0.5,        
+    ):
+        #Not implemented yet
+        print("Not implemented yet")

utils.py

@@ -1,11 +1,24 @@
 import math
 import PIL
+from PIL import Image
 import cv2
 import numpy as np
 
 from diffusers.utils import load_image
 
 def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]):
+    """
+    Draw keypoints on an image.
+
+    Args:
+    image_pil (PIL.Image): Image on which to draw the keypoints.
+    kps (list): List of keypoints to draw.
+    color_list (list): List of colors to use for drawing the keypoints.
+
+    Returns:
+    PIL.Image: Image with keypoints drawn on it.
+    """
+    
     stickwidth = 4
     limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
     kps = np.array(kps)
@@ -41,8 +54,20 @@ def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255),
 
 
 def load_and_resize_image(image_path, max_width, max_height, maintain_aspect_ratio=True):
+    """
+    Load and resize an image to the specified dimensions.
+    
+    Args:
+    image_path (str): Path to the image file.
+    max_width (int): Maximum width of the resized image.
+    max_height (int): Maximum height of the resized image.
+    maintain_aspect_ratio (bool): Whether to maintain the aspect ratio of the image.
+    
+    Returns:
+    PIL.Image: Resized image.
+    """
+
     # Open the image
-    # image = Image.open(image_path)
     image = load_image(image_path)
 
     # Get the current width and height of the image
@@ -73,7 +98,6 @@ def load_and_resize_image(image_path, max_width, max_height, maintain_aspect_rat
 
     return resized_image
 
-from PIL import Image
 
 def align_images(image1, image2):
     """
@@ -97,3 +121,45 @@ def align_images(image1, image2):
         image2 = image2.crop((0, 0, new_width, new_height))
 
     return image1, image2
+
+def align_images_2(image1, image2):
+    """
+    Resize and crop the second image to match the dimensions of the first image by
+    scaling to aspect fill and then center cropping the extra parts.
+
+    Args:
+    image1 (PIL.Image): First image which will act as the reference for alignment.
+    image2 (PIL.Image): Second image to be aligned to the first image's dimensions.
+
+    Returns:
+    tuple: A tuple containing the first image and the aligned second image.
+    """
+    # Get dimensions of the first image
+    target_width, target_height = image1.size
+
+    # Calculate the aspect ratio of the second image
+    aspect_ratio = image2.width / image2.height
+
+    # Calculate dimensions to aspect fill
+    if target_width / target_height > aspect_ratio:
+        # The first image is wider relative to its height than the second image
+        fill_height = target_height
+        fill_width = int(fill_height * aspect_ratio)
+    else:
+        # The first image is taller relative to its width than the second image
+        fill_width = target_width
+        fill_height = int(fill_width / aspect_ratio)
+
+    # Resize the second image to fill dimensions
+    filled_image = image2.resize((fill_width, fill_height), Image.Resampling.LANCZOS)
+
+    # Calculate top-left corner of crop box to center crop
+    left = (fill_width - target_width) / 2
+    top = (fill_height - target_height) / 2
+    right = left + target_width
+    bottom = top + target_height
+
+    # Crop the filled image to match the size of the first image
+    cropped_image = filled_image.crop((int(left), int(top), int(right), int(bottom)))
+
+    return image1, cropped_image