feat: integrate I-JEPA manager and HF model repository loading

README.md

@@ -8,6 +8,7 @@ sdk_version: 6.13.0
 app_file: app.py
 pinned: false
 license: mit
+python_version: 3.14
 short_description: Finds meme with the most similar expression or pose
 ---
 

app.py

@@ -0,0 +1,137 @@
+import gradio as gr
+from pathlib import Path
+from PIL import Image
+import numpy as np
+
+from config import settings
+from utils.model_loading_util import load_model_from_hf
+from src.processing.embedder import EmbeddingGenerator
+from src.search.pinecone_indexer import PineconeImageSimilaritySearch
+
+
+print("Initializing model and search index...")
+model = load_model_from_hf(
+    repo_id=settings.model_repo_id,
+    device=settings.device,
+    token=settings.hf_token
+)
+
+
+generator = EmbeddingGenerator(
+    model=model,
+    device=settings.device,
+    batch_size=settings.batch_size,
+    num_workers=2,
+    layer_strategy="last_four_concat"
+)
+
+searcher = PineconeImageSimilaritySearch(
+    index_name=settings.pinecone_index_name,
+    api_key=settings.pinecone_api_key,
+    dimension=settings.embedding_dim,
+    metric="cosine",
+)
+
+
+def search_similar_images(input_img):
+    if input_img is None:
+        return None
+
+    # Convert Gradio input (can be numpy array) to PIL Image
+    if isinstance(input_img, np.ndarray):
+        input_img = Image.fromarray(input_img)
+
+    query_embedding = generator.generate_single_embedding(input_img)
+    results = searcher.search(query_embedding, k=5, return_scores=True)
+
+    gallery_items = []
+    for path, score in results:
+        img_path = Path(path)
+        if not img_path.is_absolute():
+            img_path = settings.project_root / img_path
+
+        if img_path.exists():
+            gallery_items.append((str(img_path), f"Similarity: {score:.4f}"))
+        else:
+            print(f"Warning: Image path not found: {img_path}")
+
+    return gallery_items
+
+
+custom_css = """
+.container {
+    max-width: 1000px;
+    margin: auto;
+    padding: 20px;
+}
+.header {
+    text-align: center;
+    margin-bottom: 30px;
+}
+.header h1 {
+    font-size: 2.5rem;
+    font-weight: 800;
+    background: linear-gradient(90deg, #4F46E5, #EC4899);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    margin-bottom: 10px;
+}
+.header p {
+    color: #6B7280;
+    font-size: 1.1rem;
+}
+.gradio-container {
+    background-color: #F9FAFB !important;
+}
+.gallery-container {
+    border-radius: 12px;
+    overflow: hidden;
+    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
+}
+"""
+
+with gr.Blocks() as demo:
+    with gr.Column(elem_classes="container"):
+        with gr.Column(elem_classes="header"):
+            gr.Markdown("# Meme Similarity Search")
+            gr.Markdown(
+                "Upload an image to find the top 5 most similar memes in our database.")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    label="Upload Image",
+                    type="pil",
+                    elem_id="input-img"
+                )
+                search_btn = gr.Button("Find Similar Memes", variant="primary")
+
+            with gr.Column(scale=2):
+                output_gallery = gr.Gallery(
+                    label="Top 5 Similar Memes",
+                    show_label=True,
+                    elem_id="gallery",
+                    columns=2,
+                    rows=3,
+                    object_fit="contain",
+                    height="600px"
+                )
+
+        search_btn.click(
+            fn=search_similar_images,
+            inputs=input_image,
+            outputs=output_gallery
+        )
+
+        input_image.upload(
+            fn=search_similar_images,
+            inputs=input_image,
+            outputs=output_gallery
+        )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        css=custom_css
+    )

config/__init__.py

@@ -0,0 +1 @@
+from .settings import settings

config/settings.py

@@ -0,0 +1,32 @@
+import torch
+from pathlib import Path
+from typing import Dict, Any, Optional
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic import Field
+
+
+class Settings(BaseSettings):
+    # Pinecone Settings (Loaded from .env)
+    pinecone_api_key: str
+    pinecone_index_name: str
+
+    project_root: Path = Field(default=Path(__file__).parent.parent)
+
+    # Hugging Face Settings
+    model_repo_id: str = "Unspoiled-Egg/ijepa-target-encoder-huge"
+    hf_token: Optional[str] = None
+
+    # Model & Inference Settings
+    batch_size: int = 4
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    embedding_dim: int = 5120
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore"
+    )
+
+
+# Create the singleton instance
+settings = Settings()

requirements.txt

@@ -0,0 +1,11 @@
+"gradio>=6.13.0",
+"numpy",
+"pillow",
+"pinecone",
+"pydantic>=2.13.3",
+"pydantic-settings>=2.14.0",
+"torch",
+"torchvision",
+"tqdm>=4.67.3",
+"uuid",
+"huggingface_hub",

src/models/ijepa.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from .patch_embedding import PatchEmbed
+from .transformer_block import TransformerBlock
+
+
+class IJEPATargetEncoder(nn.Module):
+    """
+    Standard ViT without classification head.
+    Processes full image and outputs patch-level representations.
+    """
+
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        norm_layer: Optional[nn.Module] = None,
+    ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        norm_layer = norm_layer or nn.LayerNorm
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        # Positional embedding (learnable)
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches, embed_dim)
+        )
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            TransformerBlock(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+            )
+            for _ in range(depth)
+        ])
+
+        self.norm = norm_layer(embed_dim, eps=1e-6)
+
+        # Initialize weights
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_all_tokens: bool = True,
+        patch_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+          x: input images (B, C, H, W)
+          return_all_tokens: If true, return all patch tokens
+          patch_indices: If provided, return only specific patch indices
+
+        Returns:
+          Patch representations (B, N, D) or (B, len(indices), D)
+        """
+        # Patch embedding
+        x = self.patch_embed(x)   # (B, N, D)
+
+        # Add positional embeddings
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # Apply transformer blocks
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.norm(x)
+
+        # Return specific patches if indices provided
+        if patch_indices is not None:
+            x = x[:, patch_indices, :]
+
+        return x
+
+    def get_layer_representations(
+        self,
+        x: torch.Tensor,
+        strategy: str = "last",
+        specific_indices: Optional[list[int]] = None,
+        patch_indices: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Extract semantic representations using different layer strategies.
+
+        Args:
+          x: Input images (B, C, H, W)
+          strategy: Strategy to extract representations
+            - 'last': final layer only (baseline)
+            - 'second-last': 2nd-to-last block output
+            - 'last_four_concat': concat of last 4 layer (B, N, 4*D)
+            - 'specific': layers at specific_indices (eg: [25,27,29,31])
+          specific_indices: Block indices to use when strategy='specific'
+          patch_indices: If provided, return only these patch posistions
+
+        Returns:
+          (B, N, D) for "last"/"second_last", (B, N, 4*D) for concat strategies
+        """
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        n_blocks = len(self.blocks)
+
+        # Determnine which block indices to capture
+        if strategy == "second_last":
+            capture_at = {n_blocks - 2}
+        elif strategy == "last_four_concat":
+            capture_at = set(range(n_blocks - 4, n_blocks))
+        elif strategy == "specific":
+            assert specific_indices is not None, "Provide specific_indices when strategy='specific'"
+            capture_at = set(specific_indices)
+        else:
+            capture_at = {n_blocks - 1}
+
+        captured = {}
+        for i, block in enumerate(self.blocks):
+            x = block(x)
+            if i in capture_at:
+                captured[i] = x.clone()
+
+        # Apply norm and pool
+        if strategy in ("last_four_concat", "specific"):
+            # Sort by layer order, normalize each, then concat along D
+            layers = [self.norm(captured[i]) for i in sorted(captured)]
+            out = torch.cat(layers, dim=-1)   # (B, N, num_layers * D)
+        elif strategy == "second_last":
+            out = self.norm(captured[n_blocks - 2])
+        else:
+            out = self.norm(x)    # x is already the last block output
+
+        if patch_indices is not None:
+            out = out[:, patch_indices, :]
+
+        return out

src/models/multi_head_attention.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1))
+        attn = attn * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v)
+        x = x.transpose(1, 2)
+        x = x.reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

src/models/multilayer_perceptron.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x

src/models/patch_embedding.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = img_size // patch_size
+        self.num_patches = self.grid_size ** 2
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        x = x.flatten(2)
+        x = x.transpose(1, 2)
+        return x

src/models/transformer_block.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+from .multi_head_attention import MultiHeadAttention
+from .multilayer_perceptron import MLP
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.attn = MultiHeadAttention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias,
+            attn_drop=attn_drop, proj_drop=drop
+        )
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = MLP(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            drop=drop
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x

src/processing/embedder.py

@@ -0,0 +1,151 @@
+import tqdm
+import torch
+import numpy as np
+import torch.nn as nn
+from PIL import Image
+from pathlib import Path
+from torchvision import transforms
+from typing import List, Union, Tuple, Optional
+from torch.utils.data import DataLoader, Dataset
+
+
+class ImageEmbeddingDataset(Dataset):
+    """Dataset for batch image embedding generation"""
+    def __init__(
+        self,
+        image_paths: List[Union[str, Path]],
+        transform=None
+    ):
+        self.image_paths = [Path(p) for p in image_paths]
+        self.transform = transform or self.default_transform()
+
+    @staticmethod
+    def default_transform():
+        # I-JEPA uses mean=05 and std=0.5 normalization
+        return transforms.Compose([
+            transforms.Resize(
+                224, interpolation=transforms.InterpolationMode.BICUBIC
+            ),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]
+            )
+        ])
+
+    def __len__(self):
+        return len(self.image_paths)
+
+    def __getitem__(self, idx):
+        img_path = self.image_paths[idx]
+        image = Image.open(img_path).convert('RGB')
+        image = self.transform(image)
+        return image, str(img_path)
+
+
+class EmbeddingGenerator:
+    """Generate embeddings for image database using batch inference."""
+    def __init__(
+        self,
+        model: nn.Module,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        batch_size: int = 4,
+        num_workers: int = 1,
+        layer_strategy: str = "second_last",
+        specific_indices: Optional[List[int]] = None,
+    ):
+        self.model = model.to(device).eval()
+        self.device = device
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.layer_strategy = layer_strategy
+        self.specific_indices = specific_indices
+
+        # Freeze model
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def _get_features(self, images: torch.Tensor) -> torch.Tensor:
+        """Central routing method - all forward calss go through here."""
+        if self.layer_strategy == "last":
+            return self.model(images)
+        return self.model.get_layer_representations(
+            images,
+            strategy=self.layer_strategy,
+            specific_indices=self.specific_indices,
+        )
+
+    @torch.no_grad()
+    def generate_embeddings(
+        self,
+        image_paths: List[Union[str, Path]],
+        return_paths: bool = True,
+        show_progress: bool = True,
+    ) -> Union[np.ndarray, Tuple[np.ndarray, List[str]]]:
+        """
+        Generate embeddings for all images.
+
+        Returns:
+            embeddings: (N, D) array of embeddings
+            paths: (optional) list of image paths
+        """
+        print("   3.1 Image Embedding Dataset...")
+        dataset = ImageEmbeddingDataset(image_paths)
+        print("   3.2 DataLoader...")
+        dataloader = DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=True,
+        )
+
+        all_embeddings = []
+        all_paths = []
+
+        print("   3.3 tqdm iterator...\n")
+        iterator = tqdm.tqdm(dataloader, desc="Generating embeddings") if show_progress else dataloader
+
+        print("   3.4 for loop...")
+        for batch_images, batch_paths in iterator:
+            batch_images = batch_images.to(self.device, non_blocking=True)
+
+            # Get embeddings: average pool patch tokens for global representation
+            features = self._get_features(batch_images)   # (B, N, D)
+            embeddings = features.mean(dim=1)     # (B, D)
+
+            # L2 normalization for cosine similarity
+            embeddings = nn.functional.normalize(embeddings, p=2, dim=1)
+
+            all_embeddings.append(embeddings.cpu().numpy())
+            all_paths.extend(batch_paths)
+
+        embeddings = np.vstack(all_embeddings)
+
+        if return_paths:
+            return embeddings, all_paths
+        return embeddings
+
+    def generate_single_embedding(
+        self, image: Union[str, Path, Image.Image, torch.Tensor]
+    ) -> np.ndarray:
+        """Generate embedding for a single image"""
+        transform = ImageEmbeddingDataset.default_transform()
+
+        if isinstance(image, (str, Path)):
+            image = Image.open(image).convert('RGB')
+
+        if isinstance(image, Image.Image):
+            image = transform(image)
+
+        if isinstance(image, torch.Tensor):
+            image = image.unsqueeze(0) if image.dim() == 3 else image
+
+        image = image.to(self.device)
+
+        with torch.no_grad():
+            features = self._get_features(image)
+            embedding = features.mean(dim=1)
+            embedding = nn.functional.normalize(embedding, p=2, dim=1)
+
+        return embedding.cpu().numpy()

src/search/pinecone_indexer.py

@@ -0,0 +1,201 @@
+import uuid
+import numpy as np
+from pathlib import Path
+from typing import List, Union, Tuple, Optional
+
+
+class PineconeImageSimilaritySearch:
+    """
+    Pinecone-backend similarity search for image embeddings. Image paths are stored as
+    Pinecone vector metadata so no local state is needed between sessions.
+    """
+
+    def __init__(
+        self,
+        index_name: str,
+        api_key: str,
+        cloud: str = "aws",
+        region: str = "us-east-1",
+        dimension: int = 768,
+        namespace: str = "",
+        metric: str = "cosine",
+        create_if_missing: bool = True,
+    ):
+        try:
+            from pinecone import Pinecone, ServerlessSpec
+        except ImportError:
+            raise ImportError("Run `pip install pinecone` first.")
+
+        self.dimension = dimension
+        self.index_type = metric
+        self.namespace = namespace
+        self._index_name = index_name
+        self.image_paths: List[str] = []
+        self.metadata: dict = {}
+
+        pc = Pinecone(api_key=api_key)
+
+        existing = [idx.name for idx in pc.list_indexes()]
+        if index_name not in existing:
+            if not create_if_missing:
+                raise ValueError(
+                    f"Index '{index_name}' not found and create_if_missing=False."
+                )
+            pc.create_index(
+                name=index_name,
+                dimension=dimension,
+                metric=metric,
+                spec=ServerlessSpec(cloud=cloud, region=region),
+            )
+            print(f"Created Pinecone index '{index_name}' ({metric}, dim={dimension})")
+        else:
+            print(f"Connected to existing Pinecone index '{index_name}'")
+
+        self.index = pc.Index(index_name)
+
+    def add_embeddings(
+        self,
+        embeddings: np.ndarray,
+        image_paths: List[str],
+        metadata: Optional[dict] = None,
+    ):
+        """
+        Upsert embeddings into Pinecone.
+
+        Args:
+            embeddings: (N, D) float32 array for L2-normalised embeddings
+            image_paths: List of N image paths (stored as Pinecone metadata)
+            metadata: Optional {int_index: dict} of extra per-image metadata
+        """
+        assert len(embeddings) == len(image_paths), "Embeddings and paths must match"
+        assert embeddings.shape[1] == self.dimension, (
+            f"Expected dim {self.dimension}, got {embeddings.shape[1]}"
+        )
+
+        embeddings = embeddings.astype("float32")
+
+        vectors = []
+        for i, (emb, path) in enumerate(zip(embeddings, image_paths)):
+            vec_id = str(uuid.uuid4())
+            meta = {"image_path": path}
+            if metadata and i in metadata:
+                meta.update(metadata[i])
+            vectors.append({"id": vec_id, "values": emb.tolist(), "metadata": meta})
+
+        # Pinecone recommends batches of <= 100
+        batch_size = 100
+        for start in range(0, len(vectors), batch_size):
+            self.index.upsert(
+                vectors=vectors[start: start + batch_size],
+                namespace=self.namespace,
+            )
+
+        self.image_paths.extend(image_paths)
+        if metadata:
+            for i, path in enumerate(image_paths):
+                if i in metadata:
+                    self.metadata[path] = metadata[i]
+
+        print(f"Upserted {len(embeddings)} vectors. "
+              f"Total (local cache): {len(self.image_paths)}")
+
+    def search(
+        self,
+        query_embedding: np.ndarray,
+        k: int = 5,
+        return_scores: bool = True,
+        filter: Optional[dict] = None,
+    ) -> Union[List[str], List[Tuple[str, float]]]:
+        """
+        Search for the k most similar images.
+
+        Args:
+            query_embedding: (1, D) or (D,) float32 array
+            k: Number of results
+            return_scores: If True, return (path, score) tuples
+            filter: Optional Pinecone metadata filter dict
+        """
+        if query_embedding.ndim == 1:
+            query_embedding = query_embedding.reshape(1, -1)
+
+        query_list = query_embedding[0].astype("float32").tolist()
+
+        kwargs = dict(
+            vector=query_list,
+            top_k=k,
+            include_metadata=True,
+            namespace=self.namespace,
+        )
+        if filter:
+            kwargs["filter"] = filter
+
+        response = self.index.query(**kwargs)
+
+        results = []
+        for match in response.matches:
+            path = match.metadata.get("image_path", match.id)
+            if return_scores:
+                results.append((path, float(match.score)))
+            else:
+                results.append(path)
+
+        return results
+
+    def batch_search(
+        self,
+        query_embeddings: np.ndarray,
+        k: int = 5,
+        filter: Optional[dict] = None,
+    ) -> List[List[Tuple[str, float]]]:
+        """Search for multiple query embeddings sequentially."""
+        return [
+            self.search(q, k=k, return_scores=True, filter=filter)
+            for q in query_embeddings
+        ]
+
+    def save(self, save_dir: Union[str, Path]):
+        """
+        Pinecone vectors are already persisted server-side.
+        This optionally saves the local image_path cache to disk so we don't
+        have to re-scane the index on startup.
+        """
+        import pickle
+
+        save_dir = Path(save_dir)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        local_state = {
+            "image_paths": self.image_paths,
+            "metadata": self.metadata,
+            "dimension": self.dimension,
+            "namespace": self.namespace,
+            "index_name": self._index_name,
+            "metric": self.index_type,
+        }
+        with open(save_dir / "pinecone_local_cache.pkl", "wb") as f:
+            pickle.dump(local_state, f)
+        print(f"Saved local state to: {save_dir}")
+
+    @classmethod
+    def load(
+        cls, save_dir: Union[str, Path], api_key: str, use_gpu: bool = False
+    ) -> "PineconeImageSimilaritySearch":
+        """
+        Load local state and return a new PineconeImageSimilaritySearch instance.
+        """
+        import pickle
+
+        save_dir = Path(save_dir)
+        with open(save_dir / "pinecone_local_cache.pkl", "rb") as f:
+            local_state = pickle.load(f)
+
+        instance = cls(
+            index_name=local_state["index_name"],
+            api_key=api_key,
+            dimension=local_state["dimension"],
+            namespace=local_state["namespace"],
+            metric=local_state["index_type"],
+            create_if_mission=False,
+        )
+        instance.image_paths = local_state["image_paths"]
+        instance.metadata = local_state["metadata"]
+        return instance

utils/model_loading_util.py

@@ -0,0 +1,65 @@
+import json
+import torch
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
+from src.models.ijepa import IJEPATargetEncoder
+
+
+@dataclass
+class ViTConfig:
+    img_size: int = 224
+    in_chans: int = 3
+    patch_size: int = 14
+    embed_dim: int = 1280
+    depth: int = 32
+    num_heads: int = 16
+    mlp_ratio: float = 4.0
+
+
+def load_model_from_hf(
+    repo_id: str,
+    device: str = "cuda",
+    token: str = None
+):
+    """
+    Downloads and loads the I-JEPA model from a Hugging Face Model Repository.
+    """
+    print(f"Fetching model files from {repo_id}...")
+
+    # 1. Download Config
+    config_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="config.json",
+        token=token
+    )
+
+    # 2. Download Weights
+    weights_path = hf_hub_download(
+        repo_id=repo_id,
+        filename="model_weights.pth",
+        token=token
+    )
+
+    # 3. Initialize Architecture from downloaded config
+    with open(config_path, 'r') as f:
+        config_dict = json.load(f)
+    config = ViTConfig(**config_dict)
+
+    model = IJEPATargetEncoder(
+        img_size=config.img_size,
+        patch_size=config.patch_size,
+        embed_dim=config.embed_dim,
+        depth=config.depth,
+        num_heads=config.num_heads,
+        mlp_ratio=config.mlp_ratio
+    )
+
+    # 4. Load Weights
+    print("Loading state dict...")
+    state_dict = torch.load(weights_path, map_location='cpu')
+    model.load_state_dict(state_dict)
+
+    model = model.to(device).eval()
+    print("Model successfully loaded from Hugging Face.")
+
+    return model