Update IP_Composer/perform_swap.py

IP_Composer/perform_swap.py

@@ -1,36 +1,62 @@
 import torch
 import numpy as np
+from typing import List, Dict, Optional
+from PIL.Image import Image as PILImage
+from IP_Adapter import IPAdapterXL
 
-def compute_dataset_embeds_svd(all_embeds, rank):
-
+def compute_dataset_embeds_svd(
+    all_embeds: np.ndarray,
+    rank: int
+) -> np.ndarray:
     # Perform SVD on the combined matrix
-    u, s, vh = np.linalg.svd(all_embeds, full_matrices=False)
+    _, _, v = np.linalg.svd(all_embeds, full_matrices=False)
 
     # Select the top `rank` singular vectors to construct the projection matrix
-    vh = vh[:rank]  # Top `rank` right singular vectors
-    projection_matrix = vh.T @ vh  # Shape: (feature_dim, feature_dim)
+    v = v[:rank]
+    projection_matrix = v.T @ v
 
     return projection_matrix
 
-def get_embedding_composition(embed, projections_data):
-    # Initialize the combined embedding with the input embed
+def get_projected_embedding(
+    embed: np.ndarray,
+    projection_matrix: np.ndarray
+) -> np.ndarray:
+    return embed @ projection_matrix
+
+def get_embedding_composition(
+    embed: np.ndarray,
+    projections_data: List[Dict[str, np.ndarray]]
+) -> np.ndarray:
     combined_embeds = embed.copy()
 
     for proj_data in projections_data:
-
-        # Add the combined projection to the result
-        combined_embeds -= embed @ proj_data["projection_matrix"]
-        combined_embeds += proj_data["embed"] @ proj_data["projection_matrix"]
+        combined_embeds -= get_projected_embedding(embed, proj_data["projection_matrix"])
+        combined_embeds += get_projected_embedding(proj_data["embed"], proj_data["projection_matrix"])
 
     return combined_embeds
 
 
-def get_modified_images_embeds_composition(embed, projections_data, ip_model, prompt=None, scale=1.0, num_samples=3, seed=420, num_inference_steps=50):
-    
+def get_modified_images_embeds_composition(
+    embed: np.ndarray,
+    projections_data: List[Dict[str, np.ndarray]],
+    ip_model: IPAdapterXL,
+    prompt: Optional[str] = None,
+    scale: float = 1.0,
+    num_samples: int = 3,
+    seed: int = 420
+) -> List[PILImage]:
     final_embeds = get_embedding_composition(embed, projections_data)
     clip_embeds = torch.from_numpy(final_embeds)
 
-    images = ip_model.generate(clip_image_embeds=clip_embeds, prompt=prompt, num_samples=num_samples, num_inference_steps=num_inference_steps, seed=seed, guidance_scale=7.5, scale=scale)
+    images: List[PILImage] = ip_model.generate(
+    clip_image_embeds=clip_embeds,
+    prompt=prompt,
+    num_samples=num_samples,
+    num_inference_steps=50,
+    seed=seed,
+    guidance_scale=7.5,
+    scale=scale
+    )
     return images