Update app.py

app.py

@@ -1,154 +1,178 @@
-import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
+# ----------------------------------------------------------------------------
+# Copyright (c) 2024 Amar Ali-bey
+#
+# OpenVPRLab: https://github.com/amaralibey/nanoCLIP
+#
+# Licensed under the MIT License. See LICENSE file in the project root.
+# ----------------------------------------------------------------------------
 
-    generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+from pathlib import Path
+from typing import List, Tuple, Optional
 
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+import torch
+import torch.nn.functional as F
+import faiss
+from transformers import AutoTokenizer
+import gradio as gr
 
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
+from src.models import TextEncoder
+from deployment.load_album import AlbumDataset
+
+class ImageSearchEngine:
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        output_dim: int = 64,
+        gallery_folder: str = "photos",
+        device: str = 'cpu'
+    ):
+        if device == 'cuda' and not torch.cuda.is_available():
+            print("CUDA is not available. Using CPU instead.")
+            device = 'cpu'
+        self.device = torch.device(device)
+        self.setup_model(model_name, output_dim)
+        self.setup_gallery(gallery_folder)
+        
+    def setup_model(self, model_name: str, output_dim: int) -> None:
+        """Initialize and load the text encoder model."""
+        self.txt_encoder = TextEncoder(
+            output_dim=output_dim,
+            lang_model=model_name
+        ).to(self.device)
+        
+        # Load the pre-trained weights for the text encoder
+        # 
+        weights_path = Path(__file__).parent.resolve() / 'txt_encoder_state_dict.pth'
+        # check if the weights file exists
+        if not weights_path.exists():
+            raise FileNotFoundError(f"Text encoder weights not found: {weights_path}, make sure to run the create_index.py script.")
+        weights = torch.load(weights_path, map_location=self.device, weights_only=True)
+        self.txt_encoder.load_state_dict(weights)
+        self.txt_encoder.eval()
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+    def setup_gallery(self, gallery_folder: str) -> None:
+        """Setup the image gallery and FAISS index."""
+        gallery_path = Path(__file__).parent.resolve() / f'gallery/{gallery_folder}'
+        # check if the gallery folder exists
+        if not gallery_path.exists():
+            raise FileNotFoundError(f"Album folder {gallery_path} not found")
+        # we use the AlbumDataset class to load the image paths (we won't load the images themselves)
+        # this is more efficient than loading the images directly, because Gradio will load them 
+        # given the paths returned by the search method.
+        self.dataset = AlbumDataset(gallery_path, transform=None)
+        
+        # Load the FAISS index
+        # the index file should be in the same folder as the gallery 
+        # and has the same name as the folder being indexed 
+        index_path = gallery_path.parent / f"{gallery_folder}.faiss"
+        self.index = faiss.read_index(index_path.as_posix())
+        
+    @torch.no_grad()
+    def encode_query(self, query_text: str) -> torch.Tensor:
+        """Encode the text query into embeddings."""
+        inputs = self.tokenizer(query_text, truncation=True, return_tensors="pt")
+        inputs = inputs['input_ids'].to(self.device)
+        
+        embedding = self.txt_encoder(inputs)
+        embedding = F.normalize(embedding, p=2, dim=1)
+        return embedding.cpu()
+    
+    def search(self, query_text: str, k: int = 10) -> List[Tuple[str, Optional[str]]]:
+        """Search for images matching the query text."""
+        if len(query_text) < 3: # avoid searching for very short queries
+            return []
+            
+        query_embedding = self.encode_query(query_text)
+        dist, indices = self.index.search(query_embedding, k)
+        # you can filter results according to a threshold on the distance
+        return [(self.dataset.imgs[idx], None) for idx in indices[0]]
+
+class GalleryUI:
+    def __init__(self, search_engine: ImageSearchEngine):
+        self.search_engine = search_engine
+        self.css_path = Path(__file__).parent / 'style.css'
+        
+    def load_css(self) -> str:
+        """Load CSS styles from file."""
+        with open(self.css_path) as f:
+            return f.read()
+            
+    def create_interface(self) -> gr.Blocks:
+        """Create the Gradio interface."""
+        with gr.Blocks(css=self.load_css(), theme=gr.themes.Soft(text_size='lg')) as demo:
+            with gr.Column(elem_classes="container"):
+                self._create_header()
+                self._create_search_section()
+                self._create_footer()
+                
+            self._setup_callbacks(demo)
+            return demo
+    
+    def _create_header(self) -> None:
+        """Create the header section."""
+        with gr.Column(elem_classes="header"):
+            gr.Markdown("# Gallery Search")
+            gr.Markdown("Search through your collection of photos with AI")
+    
+    def _create_search_section(self) -> None:
+        """Create the search interface section."""
+        with gr.Column():
+            self.query_text = gr.Textbox(
+                placeholder="Example: Riding my horse",
+                label="Search Query",
+                elem_classes="search-input",
+                autofocus=True,
+                container=False
             )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+            
+            with gr.Row(visible=False): # this is hidden for now, but you can show it if you want
+                self.number_of_results = gr.Dropdown(
+                    choices=[4,8,12,16,24,30],
+                    value=30,
+                    label="Results per page",
+                    elem_classes="dropdown"
+                )
+                
+            self.gallery = gr.Gallery(
+                label="Search Results",
+                columns=3,
+                object_fit="cover",
+                elem_classes="gallery",
+                container=False,
             )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+    
+    def _create_footer(self) -> None:
+        """Create the footer section."""
+        with gr.Column(elem_classes="footer"):
+            gr.Markdown(
+                """Created by [Amar Ali-bey](https://amaralibey.github.io) | 
+                [View on GitHub](https://github.com/amaralibey/nanoCLIP)"""
             )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
-
+    
+    def _setup_callbacks(self, demo: gr.Blocks) -> None:
+        """Setup the interface callbacks."""
+        self.query_text.change(
+            self.search_engine.search,
+            inputs=[self.query_text, self.number_of_results],
+            outputs=self.gallery,
+            show_progress='hidden',
+        )
+        
+        self.number_of_results.change(
+            self.search_engine.search,
+            inputs=[self.query_text, self.number_of_results],
+            outputs=self.gallery
+        )
+   
 if __name__ == "__main__":
+     search_engine = ImageSearchEngine(
+        model_name = "sentence-transformers/all-MiniLM-L6-v2",
+        output_dim = 64,
+        gallery_folder = "photos",
+    )
+    ui = GalleryUI(search_engine)
+    demo = ui.create_interface()
     demo.launch()