Refactor core logic: move and modularize all latent space, sampling, and utility code into faceforge_core/

faceforge_core/attribute_directions.py

@@ -0,0 +1,35 @@
+import numpy as np
+from typing import Tuple, List, Optional
+from sklearn.decomposition import PCA
+from sklearn.linear_model import LogisticRegression
+
+class LatentDirectionFinder:
+    """
+    Provides methods to discover semantic directions in latent space using PCA or classifier-based approaches.
+    """
+    def __init__(self, latent_vectors: np.ndarray):
+        """
+        :param latent_vectors: Array of shape (N, D) where N is the number of samples and D is the latent dimension.
+        """
+        self.latent_vectors = latent_vectors
+
+    def pca_direction(self, n_components: int = 10) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Perform PCA on the latent vectors to find principal directions.
+        :return: (components, explained_variance)
+        """
+        pca = PCA(n_components=n_components)
+        pca.fit(self.latent_vectors)
+        return pca.components_, pca.explained_variance_ratio_
+
+    def classifier_direction(self, labels: List[int]) -> np.ndarray:
+        """
+        Fit a linear classifier to find a direction separating two classes in latent space.
+        :param labels: List of 0/1 labels for each latent vector.
+        :return: Normalized direction vector (D,)
+        """
+        clf = LogisticRegression()
+        clf.fit(self.latent_vectors, labels)
+        direction = clf.coef_[0]
+        direction = direction / np.linalg.norm(direction)
+        return direction 

faceforge_core/custom_loss.py

@@ -0,0 +1,26 @@
+import numpy as np
+import torch
+from typing import Callable
+
+def attribute_preserving_loss(
+    generated: torch.Tensor,
+    original: torch.Tensor,
+    attr_predictor: Callable[[torch.Tensor], torch.Tensor],
+    y_target: torch.Tensor,
+    lambda_pred: float = 1.0,
+    lambda_recon: float = 1.0
+) -> torch.Tensor:
+    """
+    Custom loss enforcing attribute fidelity and identity preservation.
+    L_attr(G(z + alpha d)) = lambda_pred * ||f_attr(G(.)) - y_target||^2 + lambda_recon * ||G(z + alpha d) - G(z)||^2
+    :param generated: Generated image tensor (B, ...)
+    :param original: Original image tensor (B, ...)
+    :param attr_predictor: Function mapping image tensor to attribute prediction
+    :param y_target: Target attribute value tensor (B, ...)
+    :param lambda_pred: Weight for attribute prediction loss
+    :param lambda_recon: Weight for reconstruction loss
+    :return: Scalar loss tensor
+    """
+    pred_loss = torch.nn.functional.mse_loss(attr_predictor(generated), y_target)
+    recon_loss = torch.nn.functional.mse_loss(generated, original)
+    return lambda_pred * pred_loss + lambda_recon * recon_loss 

faceforge_core/latent_explorer.py

@@ -0,0 +1,71 @@
+import numpy as np
+from typing import List, Optional, Tuple
+
+class LatentPoint:
+    """
+    Represents a point in latent space with an associated prompt and encoding.
+    """
+    def __init__(self, text: str, encoding: Optional[np.ndarray], xy_pos: Optional[Tuple[float, float]] = None):
+        self.text = text
+        self.encoding = encoding
+        self.xy_pos = xy_pos if xy_pos is not None else (0.0, 0.0)
+
+    def move(self, new_xy_pos: Tuple[float, float]):
+        self.xy_pos = new_xy_pos
+
+class LatentSpaceExplorer:
+    """
+    Core logic for managing points in latent space and sampling new points.
+    """
+    def __init__(self):
+        self.points: List[LatentPoint] = []
+        self.selected_point_idx: Optional[int] = None
+
+    def add_point(self, text: str, encoding: Optional[np.ndarray], xy_pos: Optional[Tuple[float, float]] = None):
+        self.points.append(LatentPoint(text, encoding, xy_pos))
+
+    def delete_point(self, idx: int):
+        if 0 <= idx < len(self.points):
+            del self.points[idx]
+
+    def modify_point(self, idx: int, new_text: str, new_encoding: Optional[np.ndarray]):
+        if 0 <= idx < len(self.points):
+            self.points[idx].text = new_text
+            self.points[idx].encoding = new_encoding
+
+    def get_encodings(self) -> List[Optional[np.ndarray]]:
+        return [p.encoding for p in self.points]
+
+    def get_prompts(self) -> List[str]:
+        return [p.text for p in self.points]
+
+    def get_positions(self) -> np.ndarray:
+        return np.array([p.xy_pos for p in self.points])
+
+    def sample_encoding(self, point: Tuple[float, float], mode: str = "distance") -> Optional[np.ndarray]:
+        """
+        Sample a new encoding based on the given point and mode.
+        """
+        encodings = self.get_encodings()
+        positions = self.get_positions()
+        if not encodings or len(encodings) == 0:
+            return None
+        if mode == "distance":
+            dists = np.linalg.norm(positions - np.array(point), axis=1)
+            coefs = 1.0 / (1.0 + dists ** 2)
+        elif mode == "circle":
+            point_vec = np.array(point)
+            positions_vec = positions
+            coefs = np.dot(positions_vec, point_vec)
+        else:
+            raise ValueError(f"Unknown sampling mode: {mode}")
+        coefs = coefs / np.sum(coefs)
+        # Weighted sum of encodings
+        result = None
+        for coef, enc in zip(coefs, encodings):
+            if enc is not None:
+                if result is None:
+                    result = coef * enc
+                else:
+                    result += coef * enc
+        return result