Initial Commit

Dockerfile

@@ -0,0 +1,29 @@
+# Dockerfile
+
+# Start with a base image that has Python and NVIDIA's CUDA drivers
+FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
+
+# Set up the environment
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+# Install Python and pip
+RUN apt-get update && apt-get install -y \
+    python3.10 python3-pip \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Copy and install Python dependencies
+COPY requirements.txt .
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+# Copy your main application code
+COPY main.py .
+
+# Expose the port your API will run on
+EXPOSE 8000
+
+# The command to start the Uvicorn server when the container runs
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

README.md

@@ -1,11 +1,9 @@
 ---
-title: Virtual Staging Api
-emoji: 📊
-colorFrom: blue
+title: Virtual Staging API
+emoji: 🛋️
+colorFrom: pink
 colorTo: blue
 sdk: docker
-pinned: false
-license: mit
+app_port: 8000
+hardware: nvidia-t4-small
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

main.py

@@ -0,0 +1,114 @@
+# main.py
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+import base64
+import io
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+from transformers import pipeline
+from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultepScheduler
+
+# --- API Data Models ---
+class StagingRequest(BaseModel):
+    image_base64: str
+    prompt: str
+    negative_prompt: str = "blurry, low quality, unrealistic, distorted, ugly, watermark, text, messy, deformed, extra windows, extra doors"
+    seed: int = 1234
+
+# --- Global State & Model Loading ---
+models = {}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # STARTUP: Load all models
+    print("🚀 Server starting up: Loading AI models...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+    models['segmentation_pipeline'] = pipeline("image-segmentation", model="Intel/dpt-large-ade", device=device)
+    models['depth_estimator'] = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas", device=device)
+    
+    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-depth", torch_dtype=torch_dtype)
+    
+    models['inpainting_pipe'] = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        controlnet=controlnet,
+        torch_dtype=torch_dtype,
+        safety_checker=None
+    ).to(device)
+    models['inpainting_pipe'].scheduler = UniPCMultistepScheduler.from_config(models['inpainting_pipe'].scheduler.config)
+    
+    print("✅ All models loaded and ready.")
+    yield
+    # SHUTDOWN: Clean up
+    print("⚡ Server shutting down.")
+    models.clear()
+
+app = FastAPI(lifespan=lifespan)
+
+# --- Helper Functions (Core Logic) ---
+def create_precise_mask(image_pil: Image.Image) -> Image.Image:
+    segments = models['segmentation_pipeline'](image_pil)
+    W, H = image_pil.size
+    inclusion_mask_np = np.zeros((H, W), dtype=np.uint8)
+    exclusion_mask_np = np.zeros((H, W), dtype=np.uint8)
+    inclusion_labels = {"wall", "floor", "ceiling"}
+    base_exclusion_labels = {"door", "window", "windowpane", "window blind"}
+    insert_labels = {"painting", "picture", "shelf", "showcase", "cabinet", "mirror", "television", "radiator"}
+    walls, inserts = [], []
+    for segment in segments:
+        label, mask = segment['label'], np.array(segment['mask'])
+        if label in inclusion_labels:
+            inclusion_mask_np = np.maximum(inclusion_mask_np, mask)
+            if label == "wall": walls.append(mask)
+        if label in base_exclusion_labels:
+            exclusion_mask_np = np.maximum(exclusion_mask_np, mask)
+        if label in insert_labels:
+            inserts.append(mask)
+    for insert_mask in inserts:
+        for wall_mask in walls:
+            if np.all((wall_mask >= insert_mask)[insert_mask > 0]):
+                exclusion_mask_np = np.maximum(exclusion_mask_np, insert_mask)
+                break
+    raw_mask_np = np.copy(inclusion_mask_np); raw_mask_np[exclusion_mask_np > 0] = 0
+    mask_filled_np = cv2.morphologyEx(raw_mask_np, cv2.MORPH_CLOSE, np.ones((10,10),np.uint8))
+    return Image.fromarray(mask_filled_np)
+
+def generate_depth_map(image_pil: Image.Image) -> Image.Image:
+    predicted_depth = models['depth_estimator'](image_pil)['predicted_depth']
+    depth_map_np = predicted_depth.cpu().numpy()
+    depth_map_np = (depth_map_np - depth_map_np.min()) / (depth_map_np.max() - depth_map_np.min()) * 255.0
+    depth_map_np = depth_map_np.astype(np.uint8)
+    return Image.fromarray(np.concatenate([depth_map_np[..., None]] * 3, axis=-1))
+
+# --- API Endpoints ---
+@app.get("/")
+def read_root():
+    return {"status": "Virtual Staging API is running."}
+
+@app.post("/furnish-room/")
+async def furnish_room(request: StagingRequest):
+    try:
+        image_bytes = base64.b64decode(request.image_base64)
+        init_image_pil = Image.open(io.BytesIO(image_bytes)).convert("RGB").resize((512, 512))
+        
+        mask_image_pil = create_precise_mask(init_image_pil)
+        control_image_pil = generate_depth_map(init_image_pil)
+        
+        generator = torch.Generator(device="cuda").manual_seed(request.seed)
+        final_image = models['inpainting_pipe'](
+            prompt=request.prompt, negative_prompt=request.negative_prompt, image=init_image_pil,
+            mask_image=mask_image_pil, control_image=control_image_pil,
+            num_inference_steps=30, guidance_scale=8.0, generator=generator,
+        ).images[0]
+        
+        buffered = io.BytesIO()
+        final_image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return {"result_image_base64": img_str}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -0,0 +1,12 @@
+# requirements.txt
+diffusers
+transformers
+accelerate
+torch
+matplotlib
+timm
+scipy
+opencv-python-headless
+fastapi
+uvicorn[standard]
+python-multipart