Update app.py

app.py

@@ -1,4 +1,3 @@
-# %%writefile app.py
 import os
 import torch
 from PIL import Image, ImageOps
@@ -10,20 +9,25 @@ import sys
 import traceback
 from datetime import datetime
 
-APP_ROOT = "/kaggle/working/virtual-staging"
+APP_ROOT = "." 
 OUTPUT_DIR = os.path.join(APP_ROOT, "outputs")
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 print(f"--- Output directory set to: {OUTPUT_DIR} ---")
+
 GROUNDING_DINO_LOCAL_PATH = os.path.join(APP_ROOT, "groundingdino_local")
 if os.path.exists(GROUNDING_DINO_LOCAL_PATH) and GROUNDING_DINO_LOCAL_PATH not in sys.path:
     sys.path.insert(0, GROUNDING_DINO_LOCAL_PATH)
     print(f"✅ Added vendorized GroundingDINO to PYTHONPATH: {GROUNDING_DINO_LOCAL_PATH}")
+
 from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler
 from transformers import pipeline as hf_pipeline
 try:
     from groundingdino.util.inference import load_model as load_gdino_model, predict as predict_gdino
     import groundingdino.datasets.transforms as T
-except ImportError as e: raise e
+except ImportError as e: 
+    print("Could not import GroundingDINO. Make sure the 'groundingdino_local' directory is in your repository.")
+    raise e
+
 HF_USERNAME = "Nightfury16"
 BASE_SD_MODEL = "runwayml/stable-diffusion-v1-5"
 CONTROLNET_INPAINT_REPO = f"{HF_USERNAME}/virtual-staging-controlnet"
@@ -39,6 +43,7 @@ def box_cxcywh_to_xyxy(x: torch.Tensor, width: int, height: int) -> torch.Tensor
     if x.nelement() == 0: return x
     x_c, y_c, w, h = x.unbind(1); b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)]
     b = torch.stack(b, dim=1); b[:, [0, 2]] *= width; b[:, [1, 3]] *= height; return b
+
 def resize_and_pad(image: Image.Image, target_size: tuple[int, int], background_color: tuple[int, int, int] = (0, 0, 0)) -> tuple[Image.Image, tuple[int, int, int, int]]:
     original_width, original_height = image.size; target_width, target_height = target_size
     ratio_w, ratio_h = target_width / original_width, target_height / original_height
@@ -49,6 +54,7 @@ def resize_and_pad(image: Image.Image, target_size: tuple[int, int], background_
     paste_x, paste_y = (target_width - new_width) // 2, (target_height - new_height) // 2
     new_image.paste(image, (paste_x, paste_y)); crop_box = (paste_x, paste_y, paste_x + new_width, paste_y + new_height)
     return new_image, crop_box
+
 class SAMModel:
     def __init__(self, device: str = 'cuda:0'): self.device, self.model = device, None
     def load(self, model_path: str = SAM_CHECKPOINT):
@@ -61,6 +67,7 @@ class SAMModel:
         final_mask = np.zeros((image.height, image.width), dtype=np.uint8)
         for mask_data in results[0].masks.data: final_mask = np.maximum(final_mask, mask_data.cpu().numpy().astype(np.uint8) * 255)
         return final_mask
+
 class DinoSamGrounding:
     def __init__(self, device: str = 'cuda:0'):
         if predict_gdino is None: raise ImportError("GroundingDINO not accessible.")
@@ -89,7 +96,6 @@ try:
 except Exception as e:
     print(f"FATAL ERROR during model loading: {e}"); global_models["loading_error"] = str(e)
 
-
 def run_virtual_staging(
     input_image: Image.Image, prompt: str, negative_prompt: str, use_canny: bool, use_depth: bool, use_lora: bool, seed: int, progress=gr.Progress()
 ):
@@ -138,6 +144,7 @@ def run_virtual_staging(
         error_message = traceback.format_exc(); print(f"!!! AN ERROR OCCURRED !!!\n{error_message}")
         raise gr.Error(f"An error occurred: {e}")
 
+
 with gr.Blocks(css="footer {display: none !important}") as demo:
     gr.Markdown("# Virtual Staging AI")
     gr.Markdown("All models are pre-loaded. Configure your generation and click 'Generate Staging'.")
@@ -171,4 +178,4 @@ with gr.Blocks(css="footer {display: none !important}") as demo:
         inputs=[input_image, prompt, negative_prompt, use_canny, use_depth, use_lora, seed_input]
     )
 
-demo.queue().launch(debug=True, share=True)
+demo.queue().launch()