Update app.py

app.py

@@ -4,8 +4,9 @@ import numpy as np
 import torch
 from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
 from sklearn.metrics.pairwise import cosine_similarity
+from PIL import Image
 import warnings
 warnings.filterwarnings('ignore')
 
@@ -120,6 +121,21 @@ button.primary:hover {
     box-shadow: 0 0 15px rgba(255, 215, 0, 0.2) !important;
     transform: translateY(-2px) !important;
 }
+
+.edit-btn {
+    background: rgba(20, 20, 20, 0.8) !important;
+    border: 1px solid rgba(255, 215, 0, 0.3) !important;
+    color: #ffd700 !important;
+    font-size: 0.9rem !important;
+    margin-top: 5px !important;
+    transition: all 0.3s ease !important;
+}
+
+.edit-btn:hover {
+    background: rgba(255, 215, 0, 0.15) !important;
+    border: 1px solid rgba(255, 215, 0, 0.8) !important;
+    box-shadow: 0 0 10px rgba(255, 215, 0, 0.3) !important;
+}
 """
 
 # --- 1. Load Dataset & Embeddings ---
@@ -143,6 +159,11 @@ pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",
 pipe = pipe.to(device)
 pipe.enable_attention_slicing()
 
+print("Loading Stable Diffusion Inpainting Model...")
+inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float32)
+inpaint_pipe = inpaint_pipe.to(device)
+inpaint_pipe.enable_attention_slicing()
+
 # --- 3. Core Engine Logic ---
 
 def get_recommendations_from_vector(user_vector):
@@ -197,6 +218,43 @@ def generate_and_recommend(prompt):
     
     return generated_image, rec1, score1, rec2, score2, rec3, score3
 
+def process_magic_edit(edit_dict, prompt):
+    if not edit_dict or not prompt:
+        return None, None, "", None, "", None, ""
+    
+    try:
+        # Extract the background and mask drawn by the user
+        bg = edit_dict["background"].convert("RGB")
+        if len(edit_dict["layers"]) > 0:
+            mask = edit_dict["layers"][0].split()[-1] # Extract the alpha channel as the mask
+        else:
+            mask = Image.new("L", bg.size, 0) # Blank mask fallback
+            
+        bg = bg.resize((512, 512))
+        mask = mask.resize((512, 512))
+        
+        # Apply Inpainting
+        edited_image = inpaint_pipe(prompt=prompt, image=bg, mask_image=mask, num_inference_steps=20).images[0]
+        
+        # Find new recommendations based on the edited image
+        inputs = processor(images=edited_image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            features = clip_model.get_image_features(**inputs)
+            if not isinstance(features, torch.Tensor):
+                features = features.pooler_output if hasattr(features, 'pooler_output') else features[0]
+                
+        user_vector = features.cpu().numpy().flatten().reshape(1, -1)
+        rec1, score1, rec2, score2, rec3, score3 = get_recommendations_from_vector(user_vector)
+        
+        return edited_image, rec1, score1, rec2, score2, rec3, score3
+    except Exception as e:
+        print(f"Error during Magic Edit: {e}")
+        return None, None, "", None, "", None, ""
+
+def send_to_magic_edit(img):
+    """Helper function to transfer an image to the Magic Edit tab and switch focus to it."""
+    return img, gr.update(selected="tab_magic")
+
 # --- 4. Gradio User Interface (Premium UI) ---
 
 custom_theme = gr.themes.Monochrome(
@@ -205,7 +263,7 @@ custom_theme = gr.themes.Monochrome(
     font=[gr.themes.GoogleFont("Montserrat"), "ui-sans-serif", "system-ui", "sans-serif"]
 )
 
-with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
+with gr.Blocks(title="Visionary | AI Interior Design", css=custom_css, theme=custom_theme) as demo:
     gr.HTML("""
     <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; margin-top: 20px;">
         <img src="https://huggingface.co/spaces/matanzig/Interior-Design-GenAI/resolve/main/viss1.jpeg" 
@@ -216,10 +274,10 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
     </div>
     """)
     
-    with gr.Tabs():
+    with gr.Tabs() as main_tabs:
         
         # --- TAB 1: Classic Search by Image ---
-        with gr.TabItem("🖼️ Search by Image"):
+        with gr.TabItem("🖼️ Search by Image", id="tab1"):
             gr.Markdown("Upload an inspiration photo to instantly discover visually and stylistically similar rooms from our curated catalog.")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -231,12 +289,15 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
                         with gr.Column():
                             img_rec1 = gr.Image(label="Top Match", elem_classes="custom-card")
                             img_score1 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            img_edit_btn1 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             img_rec2 = gr.Image(label="2nd Match", elem_classes="custom-card")
                             img_score2 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            img_edit_btn2 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             img_rec3 = gr.Image(label="3rd Match", elem_classes="custom-card")
                             img_score3 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            img_edit_btn3 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                             
             img_submit_btn.click(
                 fn=search_by_image, inputs=[image_input],
@@ -244,7 +305,7 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
             )
 
         # --- TAB 2: Fast Text Search (CLIP Multi-modal) ---
-        with gr.TabItem("🔍 Fast Text Search"):
+        with gr.TabItem("🔍 Fast Text Search", id="tab2"):
             gr.Markdown("Describe a room in text. Our multi-modal vision engine will instantly search the catalog for matching designs.")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -265,12 +326,15 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
                         with gr.Column():
                             ft_rec1 = gr.Image(label="Top Match", elem_classes="custom-card")
                             ft_score1 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            ft_edit_btn1 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             ft_rec2 = gr.Image(label="2nd Match", elem_classes="custom-card")
                             ft_score2 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            ft_edit_btn2 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             ft_rec3 = gr.Image(label="3rd Match", elem_classes="custom-card")
                             ft_score3 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            ft_edit_btn3 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
 
             fast_txt_submit_btn.click(
                 fn=search_by_text_only, inputs=[fast_text_input],
@@ -278,7 +342,7 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
             )
 
         # --- TAB 3: Advanced GenAI Search ---
-        with gr.TabItem("✨ AI Design Studio"):
+        with gr.TabItem("✨ AI Design Studio", id="tab3"):
             gr.Markdown("Describe your ideal space. Our Generative AI will draft a concept from scratch, and then find the closest real-world equivalents.")
             with gr.Row():
                 with gr.Column(scale=1):
@@ -291,21 +355,61 @@ with gr.Blocks(css=custom_css, theme=custom_theme) as demo:
                         with gr.Column():
                             txt_rec1 = gr.Image(label="Top Match", elem_classes="custom-card")
                             txt_score1 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            txt_edit_btn1 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             txt_rec2 = gr.Image(label="2nd Match", elem_classes="custom-card")
                             txt_score2 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            txt_edit_btn2 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
                         with gr.Column():
                             txt_rec3 = gr.Image(label="3rd Match", elem_classes="custom-card")
                             txt_score3 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                            txt_edit_btn3 = gr.Button("🪄 Edit Match", elem_classes="edit-btn")
 
             txt_submit_btn.click(
                 fn=generate_and_recommend, inputs=[text_input],
                 outputs=[gen_output, txt_rec1, txt_score1, txt_rec2, txt_score2, txt_rec3, txt_score3]
             )
+            
+        # --- TAB 4: Magic Edit (Inpainting) ---
+        with gr.TabItem("🎨 Magic Edit", id="tab_magic"):
+            gr.Markdown("Send a match here, paint over an object you want to change, and describe what you want instead.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    magic_editor = gr.ImageEditor(label="1. Paint over object", type="pil", elem_classes="custom-card", brush=gr.Brush(colors=["#FFFFFF"]))
+                    magic_prompt = gr.Textbox(label="2. What do you want instead?", placeholder="e.g., A modern yellow leather sofa", lines=2)
+                    magic_submit_btn = gr.Button("Apply Magic Edit & Match", variant="primary")
+                    magic_output = gr.Image(label="Magic Result", type="pil", elem_classes="custom-card")
+                    
+                with gr.Column(scale=2):
+                    with gr.Row():
+                        with gr.Column():
+                            magic_rec1 = gr.Image(label="Top Match", elem_classes="custom-card")
+                            magic_score1 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                        with gr.Column():
+                            magic_rec2 = gr.Image(label="2nd Match", elem_classes="custom-card")
+                            magic_score2 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
+                        with gr.Column():
+                            magic_rec3 = gr.Image(label="3rd Match", elem_classes="custom-card")
+                            magic_score3 = gr.Textbox(label="Confidence", interactive=False, elem_classes="custom-score")
 
-        # --- TAB 4: Presentation Video ---
-        with gr.TabItem("🎥 Presentation Video"):
+            magic_submit_btn.click(
+                fn=process_magic_edit, inputs=[magic_editor, magic_prompt],
+                outputs=[magic_output, magic_rec1, magic_score1, magic_rec2, magic_score2, magic_rec3, magic_score3]
+            )
+
+        # --- TAB 5: Presentation Video ---
+        with gr.TabItem("🎥 Presentation Video", id="tab_video"):
             gr.Markdown("### 🎓 Project Presentation & Walkthrough \nWatch the video below to see a full walkthrough of the dataset, EDA, model pipeline, and the live application.")
             gr.Video(value="https://huggingface.co/spaces/matanzig/Interior-Design-GenAI/resolve/main/A3.presentation.mp4", interactive=False)
 
+    # --- Routing logic: Connect all "Edit Match" buttons to the Magic Editor ---
+    edit_buttons_and_sources = [
+        (img_edit_btn1, img_rec1), (img_edit_btn2, img_rec2), (img_edit_btn3, img_rec3),
+        (ft_edit_btn1, ft_rec1), (ft_edit_btn2, ft_rec2), (ft_edit_btn3, ft_rec3),
+        (txt_edit_btn1, txt_rec1), (txt_edit_btn2, txt_rec2), (txt_edit_btn3, txt_rec3)
+    ]
+    
+    for btn, source_img in edit_buttons_and_sources:
+        btn.click(fn=send_to_magic_edit, inputs=[source_img], outputs=[magic_editor, main_tabs])
+
 demo.launch()