app_gradio.py.bak

import gradio as gr
import torch
import numpy as np
from PIL import Image
# import spaces  # Comment out for local development

# Model imports with error handling
try:
    from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline, UniPCMultistepScheduler
    from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
    from controlnet_aux import MLSDdetector
    MODELS_AVAILABLE = True
except ImportError as e:
    print(f"Failed to load model libraries: {e}")
    MODELS_AVAILABLE = False

# Detailed room templates by style combinations
DETAILED_PROMPTS = {
    # Living Room combinations
    ("Living Room", "Modern"): "A modern living room centered around a sleek sectional sofa and glass coffee table. A contemporary dining table with minimalist chairs provides an eating area while floor lamps and LED strips create ambient lighting throughout the clean space.",
    
    ("Living Room", "Scandinavian"): "A Scandinavian living room with a cream linen sofa and light oak coffee table complemented by a wooden dining table with simple chairs. Floating shelves display ceramics and plants while a jute rug anchors the seating area.",
    
    ("Living Room", "Industrial"): "An industrial living room anchored by a vintage leather sofa and reclaimed wood coffee table. A rustic dining table with metal chairs provides seating while pipe shelving and Edison bulb fixtures complete the urban loft aesthetic.",
    
    ("Living Room", "Bohemian"): "A bohemian living room layered with colorful Persian rugs and floor cushions around a low wooden coffee table. A vintage dining table with mismatched chairs creates a dining space while macrame wall hangings and plants in woven baskets bring life to the room.",
    
    ("Living Room", "Traditional"): "A traditional living room featuring a mahogany coffee table surrounded by wingback chairs upholstered in damask fabric. A formal dining table with upholstered chairs provides elegant seating while crystal chandeliers and antique side tables complete the classic design.",
    
    ("Living Room", "Mid-Century"): "A mid-century living room showcasing an Eames lounge chair and walnut credenza. A walnut dining table with molded plastic chairs creates a dining area while tapered leg furniture and atomic-era lighting fixtures complete the retro aesthetic.",
    
    ("Living Room", "Farmhouse"): "A farmhouse living room built around a weathered wood coffee table and slipcovered sofa. A rustic dining table with wooden chairs provides family seating while shiplap walls and mason jar lighting create authentic countryside charm.",
    
    ("Living Room", "Luxury"): "A luxury living room featuring Italian leather seating around a marble-topped coffee table. An elegant marble dining table with upholstered chairs creates sophisticated dining while crystal chandeliers and gold leaf details complete the opulent design.",
    
    # Bedroom combinations
    ("Bedroom", "Modern"): "A modern bedroom centered around a platform bed with integrated nightstands and clean geometric lines. A floating vanity and built-in wardrobes maximize space while neutral colors and accent lighting create a serene sanctuary.",
    
    ("Bedroom", "Scandinavian"): "A Scandinavian bedroom featuring a light wood bed frame with white linens and chunky knit throws. Floating nightstands hold minimalist lamps while a reading nook with sheepskin-draped chair creates the perfect hygge retreat.",
    
    ("Bedroom", "Industrial"): "An industrial bedroom with exposed brick accent wall behind a wrought iron bed frame. Metal pipe clothing racks and vintage leather trunks serve as storage while Edison bulb fixtures provide warm lighting.",
    
    ("Bedroom", "Bohemian"): "A bohemian bedroom centered around a relaxed canopy bed complemented by a large macrame wall hanging. An eclectic dresser serves as a unique storage solution while an array of potted plants brings life and color to the room.",
    
    ("Bedroom", "Traditional"): "A traditional bedroom featuring an ornate four-poster bed with mahogany finish and luxurious bedding. Matching nightstands with crystal lamps flank the bed while an antique armoire provides elegant storage.",
    
    ("Bedroom", "Mid-Century"): "A mid-century bedroom showcasing a walnut platform bed with geometric headboard and tapered leg nightstands. Atomic-era lighting fixtures complement bold graphic textiles while a vintage dresser displays period accessories.",
    
    ("Bedroom", "Farmhouse"): "A farmhouse bedroom built around a weathered wood bed frame with vintage quilt and linen pillows. A distressed dresser topped with mason jar flowers adds rustic charm while barn door closets complete the countryside aesthetic.",
    
    ("Bedroom", "Luxury"): "A luxury bedroom featuring a custom upholstered bed with tufted silk velvet headboard. Crystal chandeliers illuminate marble nightstands and cashmere bedding while a velvet chaise lounge creates an elegant sitting area.",
    
    # Kitchen combinations
    ("Kitchen", "Modern"): "A modern kitchen featuring handleless cabinets in matte charcoal with quartz waterfall countertops. Stainless steel appliances integrate seamlessly while pendant lights illuminate a large island with bar seating.",
    
    ("Kitchen", "Scandinavian"): "A Scandinavian kitchen with light oak cabinets and white marble countertops creating a clean, airy feel. Open shelving displays ceramics while natural wood bar stools surround a central island with pendant lighting.",
    
    ("Kitchen", "Industrial"): "An industrial kitchen featuring concrete countertops and exposed brick walls with black metal cabinets. Stainless steel appliances complement pipe shelving while a butcher block island and Edison bulb fixtures complete the urban aesthetic.",
    
    ("Kitchen", "Bohemian"): "A bohemian kitchen mixing vintage cabinets in sage green with colorful mosaic tile backsplash. Open shelving displays pottery and plants while copper pots hang from wrought iron racks creating artistic charm.",
    
    ("Kitchen", "Traditional"): "An elegant traditional kitchen with raised panel cabinets in warm cherry finish and granite countertops. Crown molding and decorative corbels add architectural interest, while a large island with turned legs provides additional workspace. Crystal pendant lights and oil rubbed bronze fixtures complete the classic design.",
    
    ("Kitchen", "Mid-Century"): "A retro mid-century kitchen featuring flat-front cabinets in warm walnut with stainless steel countertops. Geometric tile backsplash in turquoise and white creates visual interest, while pendant lights with atomic-era design illuminate the breakfast bar. Vintage appliances and bar stools complete the period aesthetic.",
    
    ("Kitchen", "Farmhouse"): "A rustic farmhouse kitchen with shaker cabinets in cream finish and butcher block countertops. Subway tile backsplash extends to the ceiling, while a large farmhouse sink sits below window herb gardens. Vintage-style appliances, barn door pantry, and mason jar lighting create authentic country charm.",
    
    ("Kitchen", "Luxury"): "An opulent luxury kitchen featuring custom cabinets with gold hardware and Calacatta marble countertops. Professional-grade appliances hide behind matching panels, while crystal chandeliers illuminate a large marble island. Coffered ceiling, wine storage, and fresh flowers create an atmosphere of culinary elegance.",
    
    # Dining Room combinations
    ("Dining Room", "Modern"): "A sophisticated modern dining room centered around a glass-top table with sculptural metal base, surrounded by sleek upholstered chairs. Linear chandelier provides dramatic lighting, while a built-in sideboard displays contemporary art. Floor-to-ceiling windows and neutral palette create an elegant entertaining space.",
    
    ("Dining Room", "Scandinavian"): "A warm Scandinavian dining room featuring a light oak table surrounded by wishbone chairs and illuminated by a simple pendant light. White walls showcase floating shelves with ceramics, while natural textures and plants create a cozy, family-friendly atmosphere perfect for long meals and conversation.",
    
    ("Dining Room", "Industrial"): "An edgy industrial dining room with exposed brick walls and a reclaimed wood table paired with metal chairs. Vintage factory pendant lights hang from exposed beams, while a metal pipe shelving unit displays dishes and plants. Concrete floors and vintage leather seating create authentic urban character.",
    
    ("Dining Room", "Bohemian"): "A vibrant bohemian dining room mixing vintage furniture pieces around an ornate wooden table draped with colorful textiles. Eclectic chandelier illuminates mismatched chairs, while gallery walls display an array of art and mirrors. Persian rugs and plants in macrame hangers create worldly dining charm.",
    
    ("Dining Room", "Traditional"): "A formal traditional dining room featuring a mahogany pedestal table surrounded by upholstered chairs with carved details. Crystal chandelier provides elegant lighting above fine china displayed in built-in hutch. Persian rug, oil paintings, and silk drapes create sophisticated entertaining space.",
    
    ("Dining Room", "Mid-Century"): "A stylish mid-century dining room showcasing a walnut table with geometric base surrounded by iconic molded chairs. Atomic-era chandelier illuminates vintage bar cart and credenza displaying period accessories. Bold geometric art and warm wood tones define the retro modern aesthetic.",
    
    ("Dining Room", "Farmhouse"): "A cozy farmhouse dining room built around a weathered wood trestle table with bench seating and Windsor chairs. Mason jar chandelier provides rustic lighting, while a vintage hutch displays ironstone dishes. Shiplap walls, barn wood accents, and fresh flowers create countryside dining charm.",
    
    ("Dining Room", "Luxury"): "An opulent luxury dining room featuring a custom table with marble top and gold leaf base, surrounded by velvet chairs with crystal buttons. Massive chandelier illuminates fine art and antique serving pieces displayed on marble-topped sideboard. Silk wallpaper and fresh orchids create elegant entertaining atmosphere.",
    
    # Home Office combinations
    ("Home Office", "Modern"): "A sleek modern home office featuring a minimalist white desk with integrated storage and ergonomic chair. Wall-mounted shelving displays books and decor, while hidden cable management keeps technology invisible. Large windows provide natural light, and strategic task lighting ensures productive work environment.",
    
    ("Home Office", "Scandinavian"): "A bright Scandinavian home office with light wood desk and comfortable chair positioned near window for natural light. Floating shelves display plants and books, while pegboard organizes supplies and inspiration. White walls, natural textures, and cozy lighting create perfect work-from-home hygge.",
    
    ("Home Office", "Industrial"): "An inspiring industrial home office with exposed brick walls and concrete floors, featuring a reclaimed wood desk paired with vintage leather chair. Metal pipe shelving displays books and files, while Edison bulb fixtures provide warm task lighting. Vintage typewriter and factory-style accessories complete the creative workspace.",
    
    ("Home Office", "Bohemian"): "A creative bohemian home office mixing vintage furniture pieces including ornate desk and colorful textiles. Macrame wall hangings and gallery walls provide inspiration, while plants cascade from shelves and hanging planters. Layered rugs, brass accessories, and natural light create artistic work environment.",
    
    ("Home Office", "Traditional"): "A sophisticated traditional home office featuring rich mahogany desk with leather inlay and matching bookshelf filled with leather-bound volumes. Elegant table lamp provides task lighting, while Persian rug and oil paintings create refined scholarly atmosphere. Wing-back chair offers comfortable reading spot.",
    
    ("Home Office", "Mid-Century"): "A stylish mid-century home office showcasing iconic walnut desk with hairpin legs and molded chair. Atomic-era task lamp illuminates work surface, while geometric shelving displays books and period accessories. Bold artwork and warm wood tones create inspiring retro modern workspace.",
    
    ("Home Office", "Farmhouse"): "A cozy farmhouse home office built around weathered wood desk with vintage chair and mason jar storage. Shiplap walls display inspiration boards and vintage signs, while galvanized accessories organize supplies. Natural light from window and rustic chandelier create productive countryside workspace.",
    
    ("Home Office", "Luxury"): "An opulent luxury home office featuring custom built-ins in rich mahogany with gold hardware and marble accents. Executive leather chair faces elegant desk with crystal accessories, while chandelier illuminates fine art and fresh flowers. Silk drapes and Persian rug create sophisticated work environment.",
    
    # Bathroom combinations
    ("Bathroom", "Modern"): "A spa-like modern bathroom featuring floating vanity with vessel sinks and waterfall faucets. Floor-to-ceiling tiles create seamless surfaces, while frameless glass shower and freestanding tub maximize the minimalist aesthetic. Hidden LED lighting and natural stone accents create serene sanctuary.",
    
    ("Bathroom", "Scandinavian"): "A bright Scandinavian bathroom with light wood vanity and white marble countertops creating clean, airy feel. Natural textures include woven baskets for storage and wood bath caddy, while plants thrive in the humid environment. Simple fixtures and abundant natural light complete the Nordic spa aesthetic.",
    
    ("Bathroom", "Industrial"): "An edgy industrial bathroom with exposed pipes and concrete vanity topped with vessel sinks. Subway tiles and metal fixtures complement vintage mirror and Edison bulb lighting, while cast iron tub anchors the space. Raw materials and utilitarian design create authentic urban loft character.",
    
    ("Bathroom", "Bohemian"): "A luxurious bohemian bathroom mixing vintage furniture pieces as vanity with ornate mirror and colorful mosaic tiles. Macrame plant hangers and layered textiles create spa-like atmosphere, while clawfoot tub surrounded by candles offers relaxing retreat. Eclectic accessories and natural elements complete the worldly aesthetic.",
    
    ("Bathroom", "Traditional"): "An elegant traditional bathroom featuring marble countertops with undermount sinks and polished brass fixtures. Wainscoting and crown molding add architectural detail, while crystal chandelier provides luxury lighting. Clawfoot tub, Persian rug, and fresh flowers complete the sophisticated spa experience.",
    
    ("Bathroom", "Mid-Century"): "A stylish mid-century bathroom showcasing geometric tiles and walnut vanity with brass fixtures. Atomic-era mirror and lighting fixtures complement sleek lines, while sunken tub and bold accent colors define the retro modern aesthetic. Period accessories and warm wood tones complete the vintage look.",
    
    ("Bathroom", "Farmhouse"): "A cozy farmhouse bathroom with shiplap walls and weathered wood vanity topped with vessel sinks. Galvanized fixtures and mason jar lighting create rustic charm, while clawfoot tub surrounded by vintage accessories offers relaxing retreat. Natural textures and countryside elements complete the authentic aesthetic.",
    
    ("Bathroom", "Luxury"): "An opulent luxury bathroom featuring marble surfaces throughout with gold fixtures and crystal chandelier. Freestanding soaking tub faces fireplace, while double vanity offers ample storage. Heated floors, fresh orchids, and plush towels create five-star spa experience at home.",
    
    # Kids Room combinations
    ("Kids Room", "Modern"): "A sleek modern kids room featuring built-in bunk beds with integrated storage and study areas. Bright accent colors pop against white walls, while interactive technology and educational displays encourage learning. Smart storage solutions and safety features create functional space for growing children.",
    
    ("Kids Room", "Scandinavian"): "A cozy Scandinavian kids room with light wood furniture and neutral colors creating calm environment. Natural toys and books display on floating shelves, while cozy reading nook with sheepskin throw encourages quiet time. Simple design and quality materials create timeless childhood space.",
    
    ("Kids Room", "Industrial"): "A creative industrial kids room with exposed elements softened by colorful textiles and playful accessories. Metal pipe clothing rack and vintage trunks provide storage, while chalkboard walls encourage artistic expression. Edison bulb fixtures and reclaimed wood furniture create unique urban nursery.",
    
    ("Kids Room", "Bohemian"): "A whimsical bohemian kids room layered with colorful textiles and global-inspired decor. Teepee reading corner and floor cushions create cozy play areas, while macrame details and plants add natural elements. Eclectic furniture and artistic displays encourage creativity and imagination.",
    
    ("Kids Room", "Traditional"): "A classic traditional kids room featuring quality wood furniture including four-poster bed and matching dresser. Timeless patterns and colors create sophisticated yet age-appropriate space, while built-in bookcases encourage reading. Heirloom quality pieces grow with child through years.",
    
    ("Kids Room", "Mid-Century"): "A playful mid-century kids room showcasing period furniture in miniature scale with bold geometric patterns. Atomic-era accessories and retro color palette create fun vintage atmosphere, while built-in storage and study areas support growing needs. Quality design principles create lasting childhood memories.",
    
    ("Kids Room", "Farmhouse"): "A charming farmhouse kids room with weathered wood furniture and vintage accessories creating countryside appeal. Shiplap accent wall displays family photos and artwork, while galvanized storage bins organize toys. Natural materials and rustic charm create wholesome childhood environment.",
    
    ("Kids Room", "Luxury"): "An elegant luxury kids room featuring custom built-ins and high-end materials adapted for young users. Crystal chandelier and silk curtains create sophisticated atmosphere, while quality furniture and accessories ensure lasting beauty. Premium materials and thoughtful design create special childhood sanctuary.",
    
    # Master Bedroom combinations
    ("Master Bedroom", "Modern"): "A sophisticated modern master bedroom featuring king platform bed with integrated nightstands and dramatic headboard wall. Floor-to-ceiling windows with automated controls provide natural light and privacy, while walk-in closet and ensuite bathroom create luxurious retreat. Neutral palette and clean lines ensure restful environment.",
    
    ("Master Bedroom", "Scandinavian"): "A serene Scandinavian master bedroom with light wood bed frame and crisp white linens creating peaceful sanctuary. Cozy sitting area by window includes reading chair and side table, while natural textures and muted colors promote relaxation. Quality materials and simple design create timeless bedroom retreat.",
    
    ("Master Bedroom", "Industrial"): "A dramatic industrial master bedroom with exposed brick accent wall and steel beam ceiling details. Vintage leather furniture and metal accents complement platform bed, while factory-style lighting provides ambient illumination. Raw materials and urban aesthetics create unique romantic retreat.",
    
    ("Master Bedroom", "Bohemian"): "A luxurious bohemian master bedroom with canopy bed draped in flowing fabrics and surrounded by eclectic vintage furniture. Layered textiles in rich colors create cozy atmosphere, while plants and global accessories add worldly charm. Artistic details and natural elements create romantic sanctuary.",
    
    ("Master Bedroom", "Traditional"): "An elegant traditional master bedroom featuring ornate four-poster bed with luxury bedding and matching furniture suite. Crystal chandelier and silk drapes create formal atmosphere, while sitting area and fireplace add comfort. Classic design elements and quality materials create timeless romantic retreat.",
    
    ("Master Bedroom", "Mid-Century"): "A stylish mid-century master bedroom showcasing iconic furniture pieces including walnut platform bed and vintage accessories. Bold geometric patterns and warm wood tones create retro modern atmosphere, while period lighting and textiles complete the vintage aesthetic. Quality design creates lasting bedroom sanctuary.",
    
    ("Master Bedroom", "Farmhouse"): "A cozy farmhouse master bedroom built around weathered wood bed frame with vintage quilt and natural linens. Shiplap walls and rustic accessories create countryside charm, while sitting area by window offers peaceful retreat. Natural materials and authentic details create romantic rural sanctuary.",
    
    ("Master Bedroom", "Luxury"): "An opulent luxury master bedroom featuring custom upholstered bed with silk velvet headboard and crystal chandelier illumination. Marble fireplace anchors sitting area with velvet chairs, while walk-in closet and spa bathroom create five-star hotel experience. Premium materials and elegant details create ultimate romantic retreat."
}

# Room types and styles for dropdowns
ROOM_TYPES = list(set([combo[0] for combo in DETAILED_PROMPTS.keys()]))
STYLE_TYPES = list(set([combo[1] for combo in DETAILED_PROMPTS.keys()]))

# Global variables for models
pipe = None
seg_processor = None
seg_model = None
mlsd_processor = None

def load_models():
    """Load models (called once)"""
    global pipe, seg_processor, seg_model, mlsd_processor
    
    if not MODELS_AVAILABLE:
        return "❌ Model libraries not available"
    
    try:
        print("🔄 Loading AI models...")
        
        # Dual ControlNet setup
        controlnet = [
            ControlNetModel.from_pretrained(
                "BertChristiaens/controlnet-seg-room", 
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
            ),
            ControlNetModel.from_pretrained(
                "lllyasviel/sd-controlnet-mlsd", 
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
            ),
        ]
        
        # Main pipeline
        pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
            "SG161222/Realistic_Vision_V3.0_VAE",
            controlnet=controlnet,
            safety_checker=None,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        
        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
        
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")
        
        # Segmentation models
        seg_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640")
        seg_model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640")
        mlsd_processor = MLSDdetector.from_pretrained("lllyasviel/Annotators")
        
        print("✅ Models loaded successfully!")
        return "✅ Models loaded successfully!"
        
    except Exception as e:
        return f"❌ Failed to load models: {e}"

def create_full_mask(image):
    """Create full image mask to eliminate all boundaries"""
    w, h = image.size
    
    # Full white mask - allows modification of entire image
    mask = np.ones((h, w), dtype=np.uint8) * 255
    
    return Image.fromarray(mask).convert("RGB")

def create_smart_mask(image):
    """Create conservative smart mask that primarily targets floor center area"""
    import cv2
    
    # Convert PIL to numpy
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    
    # Create a very conservative mask - only target center floor area
    mask = np.zeros((h, w), dtype=np.uint8)
    
    # Define safe zone - center 60% of image, bottom 70% (floor area)
    safe_x_start = int(w * 0.2)  # 20% from left
    safe_x_end = int(w * 0.8)    # 20% from right  
    safe_y_start = int(h * 0.3)  # 30% from top (avoid ceiling)
    safe_y_end = int(h * 0.9)    # 10% from bottom
    
    # Create elliptical mask in center floor area
    center_x, center_y = w // 2, int(h * 0.65)  # Slightly lower center
    
    # Create ellipse parameters - larger size for better furniture generation
    ellipse_w = int(w * 0.40)  # 40% of width (larger area)
    ellipse_h = int(h * 0.28)  # 28% of height (larger area)
    
    # Draw ellipse mask
    Y, X = np.ogrid[:h, :w]
    ellipse_mask = ((X - center_x) / ellipse_w) ** 2 + ((Y - center_y) / ellipse_h) ** 2 <= 1
    
    # Apply gradient within ellipse - strongest at center, fade to edges
    for y in range(h):
        for x in range(w):
            if ellipse_mask[y, x]:
                # Distance from ellipse center
                dist = np.sqrt(((x - center_x) / ellipse_w) ** 2 + ((y - center_y) / ellipse_h) ** 2)
                # Gradient: stronger at center (dist=0), weaker at edges (dist=1)
                strength = max(0, 1 - dist) * 255  # Max 255 - full strength
                mask[y, x] = int(strength)
    
    # Additional safety: exclude areas that are too bright (likely windows) or too dark (likely corners)
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    
    # Mask out very bright areas (windows) and very dark areas (corners/shadows)
    bright_mask = gray > 220  # Very bright
    dark_mask = gray < 40     # Very dark
    
    # Reduced corner protection - exclude outer 5% of image only
    corner_margin = 0.05
    corner_mask = np.zeros((h, w), dtype=bool)
    corner_mask[:int(h*corner_margin), :] = True  # Top
    corner_mask[-int(h*corner_margin):, :] = True  # Bottom
    corner_mask[:, :int(w*corner_margin)] = True  # Left
    corner_mask[:, -int(w*corner_margin):] = True  # Right
    
    exclude_mask = bright_mask | dark_mask | corner_mask
    mask[exclude_mask] = 0
    
    # Apply blur for smooth transitions
    mask = cv2.GaussianBlur(mask, (21, 21), 0)
    
    # Debug: print mask statistics
    print(f"Mask stats - Min: {mask.min()}, Max: {mask.max()}, Non-zero pixels: {np.count_nonzero(mask)}")
    print(f"Mask shape: {mask.shape}, Image shape: {img_array.shape}")
    
    # Save mask for debugging (temporary)
    mask_debug = Image.fromarray(mask).convert("RGB")
    try:
        mask_debug.save("/tmp/debug_mask.png")
        print("Debug mask saved to /tmp/debug_mask.png")
    except:
        pass
    
    return mask_debug

def analyze_room_structure(image):
    """NEVER TOUCH STRUCTURE: Analyze walls, windows, ceiling - READ ONLY"""
    import cv2
    
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    
    # Detect structural elements that must NEVER be modified
    structure_mask = np.zeros((h, w), dtype=np.uint8)
    
    # 1. Wall detection (strong vertical/horizontal lines)
    edges = cv2.Canny(gray, 50, 150)
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=50, minLineLength=50, maxLineGap=10)
    
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            cv2.line(structure_mask, (x1, y1), (x2, y2), 255, 5)
    
    # 2. Window detection (very bright rectangular areas)
    bright_mask = gray > 200
    contours, _ = cv2.findContours(bright_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in contours:
        area = cv2.contourArea(contour)
        if area > 1000:  # Large bright areas = windows
            cv2.fillPoly(structure_mask, [contour], 255)
    
    # 3. Ceiling detection (top 30% of image)
    structure_mask[:int(h*0.3), :] = 255
    
    # 4. Wall edges (outer 15% of image)
    border = int(min(w, h) * 0.15)
    structure_mask[:border, :] = 255      # Top
    structure_mask[-border:, :] = 255     # Bottom
    structure_mask[:, :border] = 255      # Left
    structure_mask[:, -border:] = 255     # Right
    
    # Expand structure protection
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))
    structure_mask = cv2.dilate(structure_mask, kernel, iterations=2)
    
    return structure_mask

def detect_floor_area(image, structure_mask):
    """Physics-based floor detection with hard constraints"""
    import cv2
    
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    
    # Floor must be in bottom 70% of image (physics constraint)
    floor_region = np.zeros((h, w), dtype=np.uint8)
    floor_start_y = int(h * 0.3)  # Floor cannot be in top 30%
    
    # Detect horizontal surfaces (floor characteristics)
    gray_floor = gray[floor_start_y:, :]
    
    # Horizontal gradient analysis - floors have low vertical variation
    grad_y = cv2.Sobel(gray_floor, cv2.CV_64F, 0, 1, ksize=3)
    horizontal_surfaces = np.abs(grad_y) < 15
    
    # Consistent texture/color (floor property)
    blurred = cv2.GaussianBlur(gray_floor, (21, 21), 0)
    floor_brightness = np.median(blurred)
    consistent_areas = np.abs(blurred - floor_brightness) < 25
    
    # Combine constraints: horizontal + consistent = floor
    floor_candidates = horizontal_surfaces & consistent_areas
    floor_region[floor_start_y:, :] = floor_candidates.astype(np.uint8) * 255
    
    # HARD CONSTRAINT: Remove any overlap with structure
    floor_region[structure_mask > 0] = 0
    
    # Physics validation: floor must be connected and substantial
    contours, _ = cv2.findContours(floor_region, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        # Keep only largest floor area (physics: one continuous floor)
        largest = max(contours, key=cv2.contourArea)
        floor_region = np.zeros_like(floor_region)
        cv2.fillPoly(floor_region, [largest], 255)
    
    return floor_region

def create_furniture_placement_zones(floor_mask, image):
    """Define valid furniture placement with physics constraints"""
    import cv2
    
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    
    # Physics constraints for furniture placement
    placement_zones = np.zeros((h, w), dtype=np.uint8)
    
    if np.any(floor_mask > 0):
        # Find floor center of mass (realistic furniture placement)
        moments = cv2.moments(floor_mask)
        if moments["m00"] != 0:
            cx = int(moments["m10"] / moments["m00"])
            cy = int(moments["m01"] / moments["m00"])
        else:
            cx, cy = w//2, int(h*0.7)
        
        # Create placement zones around floor center
        # Bedroom: bed against wall, nightstands beside
        if True:  # For now, general furniture placement
            # Main furniture zone (bed, sofa, etc.)
            main_radius_x = int(w * 0.25)
            main_radius_y = int(h * 0.15)
            
            Y, X = np.ogrid[:h, :w]
            main_zone = ((X - cx) / main_radius_x) ** 2 + ((Y - cy) / main_radius_y) ** 2 <= 1
            placement_zones[main_zone & (floor_mask > 0)] = 255
        
        # Only place furniture on detected floor
        placement_zones[floor_mask == 0] = 0
    
    return placement_zones

def create_layered_furniture_mask(image):
    """LAYERED APPROACH: Allow furniture on walls but prevent structure replacement"""
    import cv2
    
    print("🎨 Creating layered furniture mask...")
    
    img_array = np.array(image)
    h, w = img_array.shape[:2]
    gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
    
    # Create large mask covering most of the room
    mask = np.zeros((h, w), dtype=np.uint8)
    
    # Main furniture area - much larger to allow wall furniture
    center_x = w // 2
    center_y = int(h * 0.6)  # 60% down from top
    
    # Large coverage area for layered furniture
    radius_x = int(w * 0.45)  # 45% of width - covers most room
    radius_y = int(h * 0.35)  # 35% of height - includes wall areas
    
    # Create large elliptical coverage
    Y, X = np.ogrid[:h, :w]
    main_ellipse = ((X - center_x) / radius_x) ** 2 + ((Y - center_y) / radius_y) ** 2 <= 1
    
    # Add additional area for wall furniture (bookcases, wall art, etc.)
    wall_coverage = np.zeros((h, w), dtype=bool)
    wall_coverage[int(h*0.2):int(h*0.9), int(w*0.1):int(w*0.9)] = True
    
    # Combine main area with wall coverage
    furniture_area = main_ellipse | wall_coverage
    
    # Create gradient mask - stronger in center, lighter near edges
    for y in range(h):
        for x in range(w):
            if furniture_area[y, x]:
                # Calculate distance from center
                dist_from_center = np.sqrt((x - center_x)**2 + (y - center_y)**2)
                max_dist = np.sqrt(radius_x**2 + radius_y**2)
                
                # Gradient: strong in center, medium at edges
                if dist_from_center <= max_dist * 0.6:
                    mask[y, x] = 255  # Full strength in center
                elif dist_from_center <= max_dist:
                    # Fade from 255 to 180 towards edges
                    strength = 255 - int((dist_from_center - max_dist * 0.6) / (max_dist * 0.4) * 75)
                    mask[y, x] = max(180, strength)
                else:
                    # Light coverage for wall areas
                    mask[y, x] = 180
    
    # Exclude only window areas (very bright)
    windows = gray > 220
    mask[windows] = 0
    
    # Exclude extreme corners (likely ceiling/wall joints)
    corner_size = int(min(w, h) * 0.05)
    mask[:corner_size, :corner_size] = 0      # Top-left
    mask[:corner_size, -corner_size:] = 0     # Top-right
    
    # Smooth transitions
    mask = cv2.GaussianBlur(mask, (5, 5), 0)
    
    # Statistics
    furniture_pixels = np.count_nonzero(mask)
    coverage = (furniture_pixels / (h * w)) * 100
    strong_pixels = np.count_nonzero(mask > 240)
    medium_pixels = np.count_nonzero((mask > 160) & (mask <= 240))
    
    print(f"🎯 Layered mask coverage:")
    print(f"  - Total coverage: {furniture_pixels} pixels ({coverage:.1f}%)")
    print(f"  - Strong areas (floor): {strong_pixels} pixels")
    print(f"  - Medium areas (walls): {medium_pixels} pixels")
    print(f"  - Center: ({center_x}, {center_y}), Size: {radius_x}x{radius_y}")
    
    return Image.fromarray(mask).convert("RGB")

# Keep the old function for compatibility
def create_precise_furniture_mask(image):
    """Legacy function - redirects to layered approach"""
    return create_layered_furniture_mask(image)

def get_prompt_preview(room_type, design_style, inpainting_mode):
    """Generate preview of prompt and negative prompt that will be used"""
    
    # Create positive prompt based on mode
    if inpainting_mode == "smart":
        # Simple, direct furniture prompt for smart mode
        if room_type == "Living Room":
            furniture_items = "sofa, coffee table, side tables, floor lamp, dining table with chairs"
        elif room_type == "Bedroom" or room_type == "Master Bedroom":
            furniture_items = "MUST INCLUDE: large bed with headboard, two nightstands with lamps, dresser or wardrobe, accent chair, area rug under bed"
        elif room_type == "Kitchen":
            furniture_items = "kitchen island, bar stools, dining table with chairs"
        elif room_type == "Dining Room":
            furniture_items = "dining table, dining chairs, sideboard"
        elif room_type == "Home Office":
            furniture_items = "desk, office chair, bookshelf, filing cabinet"
        elif room_type == "Bathroom":
            furniture_items = "vanity, mirror, storage cabinet"
        else:
            furniture_items = "appropriate furniture"
        
        positive_prompt = f"LAYER-BASED GENERATION: add {furniture_items} ONLY on floor surface, {design_style.lower()} style furniture, photorealistic, NEVER TOUCH STRUCTURE: walls ceiling windows floors are READ-only, HARD CONSTRAINT: furniture objects placed on existing floor only, PHYSICS VALIDATION: realistic furniture placement, professional interior design, ultra-realistic, perfect lighting integration"
    else:
        # Get detailed template-based prompt for full mode
        detailed_prompt = DETAILED_PROMPTS.get((room_type, design_style), 
                                             DETAILED_PROMPTS[("Living Room", "Modern")])
        positive_prompt = f"photorealistic interior design, {detailed_prompt}, keep existing windows unchanged, preserve original window placement, professionally photographed, architectural photography, natural lighting, ultra-realistic, high resolution, sharp focus, interior design magazine quality, realistic textures, realistic materials"
    
    # HARD CONSTRAINTS negative prompt
    negative_prompt = "STRUCTURE MODIFICATION FORBIDDEN: changing walls, modifying walls, different walls, new walls, changing ceiling, modifying ceiling, different ceiling, changing floor, modifying floor, different floor, new windows, extra windows, additional windows, changing windows, modifying windows, different windows, window alterations, removing windows, PHYSICS VIOLATIONS: floating furniture, unrealistic placement, furniture in walls, furniture in ceiling, structural changes, architectural changes, room alterations, upholstered walls, fabric walls, plush walls, lowres, watermark, banner, logo, contactinfo, text, deformed, blurry, blur, out of focus, surreal, ugly"
    
    return positive_prompt, negative_prompt

def post_process_blend(original, generated):
    """Post-process to reduce seam artifacts between original and generated areas"""
    from PIL import ImageFilter
    
    # Apply slight blur to reduce harsh transitions
    blended = generated.filter(ImageFilter.GaussianBlur(radius=0.8))
    
    # Blend with original in border areas for smoother transitions
    orig_array = np.array(original, dtype=np.float32)
    gen_array = np.array(blended, dtype=np.float32)
    
    h, w = orig_array.shape[:2]
    
    # Create blend mask - stronger blending near edges
    blend_mask = np.ones((h, w), dtype=np.float32)
    
    # Reduce blend strength near borders
    border_size = min(h, w) // 20
    for i in range(border_size):
        alpha = i / border_size
        blend_mask[i, :] = alpha
        blend_mask[-(i+1), :] = alpha
        blend_mask[:, i] = np.minimum(blend_mask[:, i], alpha)
        blend_mask[:, -(i+1)] = np.minimum(blend_mask[:, -(i+1)], alpha)
    
    # Apply blending
    if len(orig_array.shape) == 3:
        blend_mask = blend_mask[:, :, np.newaxis]
    
    final_array = orig_array * (1 - blend_mask) + gen_array * blend_mask
    final_array = np.clip(final_array, 0, 255).astype(np.uint8)
    
    return Image.fromarray(final_array)

# @spaces.GPU  # Comment out for local development
def design_space(input_image, room_type, design_style, inpainting_mode, num_steps, guidance_scale, strength):
    """Generate space design using ZeroGPU"""
    
    global pipe, seg_processor, seg_model, mlsd_processor
    
    if input_image is None:
        return None, "❌ Please upload an image!"
    
    if pipe is None:
        # Load models if not already loaded
        status = load_models()
        if "❌" in status:
            return None, status
    
    try:
        # Use detailed template-based prompt with photorealistic emphasis
        detailed_prompt = DETAILED_PROMPTS.get((room_type, design_style), 
                                             DETAILED_PROMPTS[("Living Room", "Modern")])
        prompt = f"photorealistic interior design, {detailed_prompt}, professionally photographed, architectural photography, natural lighting, ultra-realistic, high resolution, sharp focus, interior design magazine quality, realistic textures, realistic materials"
        prompt_type = f"{room_type} in {design_style} style"
        
        # Resize image
        orig_w, orig_h = input_image.size
        max_size = 768
        if max(orig_w, orig_h) > max_size:
            if orig_w > orig_h:
                new_w, new_h = max_size, int(max_size * orig_h / orig_w)
            else:
                new_w, new_h = int(max_size * orig_w / orig_h), max_size
        else:
            new_w, new_h = orig_w, orig_h
            
        resized_image = input_image.resize((new_w, new_h))
        
        # Simple segmentation (create basic control image)
        seg_control = resized_image.copy()
        
        # MLSD processing
        if mlsd_processor:
            mlsd_image = mlsd_processor(resized_image)
            mlsd_image = mlsd_image.resize((new_w, new_h))
        else:
            mlsd_image = resized_image.copy()
        
        # Create mask based on selected mode
        if inpainting_mode == "smart":
            # Use precise smart mask for furniture-only placement
            mask_image = create_precise_furniture_mask(resized_image)
            
            # Simple, direct furniture prompt for smart mode
            if room_type == "Living Room":
                furniture_items = "sofa, coffee table, side tables, floor lamp, dining table with chairs"
            elif room_type == "Bedroom":
                furniture_items = "bed, nightstands, dresser, chair"
            elif room_type == "Kitchen":
                furniture_items = "kitchen island, bar stools, dining table with chairs"
            elif room_type == "Dining Room":
                furniture_items = "dining table, dining chairs, sideboard"
            elif room_type == "Home Office":
                furniture_items = "desk, office chair, bookshelf, filing cabinet"  
            elif room_type == "Bathroom":
                furniture_items = "vanity, mirror, storage cabinet"
            else:
                furniture_items = "appropriate furniture"
            
            prompt = f"add {furniture_items} to this room, {design_style.lower()} style furniture, photorealistic, keep all walls ceiling floor windows unchanged, preserve all existing room structure, only add furniture objects on the floor, do not modify room architecture, professional interior design, ultra-realistic, high quality"
            print(f"Smart mode prompt: {prompt}")
        else:
            mask_image = create_full_mask(resized_image)
            prompt = f"photorealistic interior design, {detailed_prompt}, keep existing windows unchanged, preserve original window placement, professionally photographed, architectural photography, natural lighting, ultra-realistic, high resolution, sharp focus, interior design magazine quality, realistic textures, realistic materials"
            print("Using full mask mode")
        
        # Generate image with optimized strength for furniture generation
        if inpainting_mode == "smart":
            # Higher strength for better furniture generation in smart mode
            actual_strength = min(0.85, float(strength) * 1.1)  # Boost by 10%, cap at 0.85
        else:
            actual_strength = float(strength)
        
        print(f"Generation parameters:")
        print(f"  - Mode: {inpainting_mode}")
        print(f"  - Original strength: {strength}")
        print(f"  - Actual strength: {actual_strength}")
        print(f"  - Steps: {int(num_steps)}")
        print(f"  - Guidance scale: {float(guidance_scale)}")
        
        result = pipe(
            prompt=prompt,
            negative_prompt="lowres, watermark, banner, logo, watermark, contactinfo, text, deformed, blurry, blur, out of focus, out of frame, surreal, extra, ugly, upholstered walls, fabric walls, plush walls, mirror, mirrored, functional, realistic, new windows, extra windows, additional windows, changing windows, modifying windows, different windows, window alterations, removing windows, changing walls, modifying walls, different walls, new walls, changing ceiling, modifying ceiling, different ceiling, changing floor, modifying floor, different floor, structural changes, architectural changes, room alterations",
            num_inference_steps=int(num_steps),
            strength=actual_strength,
            guidance_scale=float(guidance_scale),
            image=resized_image,
            mask_image=mask_image,
            control_image=[seg_control, mlsd_image],
            controlnet_conditioning_scale=[0.4, 0.2],
            control_guidance_start=[0, 0.1],
            control_guidance_end=[0.5, 0.25],
        ).images[0]
        
        # Restore original size
        final_image = result.resize((orig_w, orig_h), Image.Resampling.LANCZOS)
        
        # No post-processing needed with full mask
        
        success_msg = f"✅ {prompt_type} completed! Generated in {int(num_steps)} steps."
        
        return final_image, success_msg
        
    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        return None, error_msg

# Gradio interface
def create_interface():
    """Create Gradio interface"""
    
    with gr.Blocks(title="Spacely AI Interior Designer", theme=gr.themes.Soft()) as demo:
        gr.HTML("<h1>🏠 Spacely AI Interior Designer</h1>")
        gr.Markdown("Upload an empty room photo and AI will design it with furniture")
        
        with gr.Row():
            with gr.Column(scale=1):
                # Input controls
                input_image = gr.Image(
                    label="Upload Empty Room Image",
                    type="pil",
                    height=300
                )
                
                room_type = gr.Dropdown(
                    choices=ROOM_TYPES,
                    value="Living Room",
                    label="Select Room Type"
                )
                
                design_style = gr.Dropdown(
                    choices=STYLE_TYPES,
                    value="Modern",
                    label="Select Design Style"
                )
                
                inpainting_mode = gr.Radio(
                    choices=[
                        ("Complete Room Redesign", "full"),
                        ("Add Furniture Only (Preserve Walls)", "smart")
                    ],
                    value="smart",
                    label="🎨 Design Mode",
                    info="Choose how to modify your image"
                )
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    num_steps = gr.Slider(
                        minimum=1, maximum=500, value=50, step=1,
                        label="Number of denoising steps"
                    )
                    guidance_scale = gr.Slider(
                        minimum=1, maximum=50, value=18, step=0.5,
                        label="Scale for classifier-free guidance"
                    )
                    strength = gr.Slider(
                        minimum=0, maximum=1, value=0.8, step=0.05,
                        label="Prompt strength for inpainting"
                    )
                
                generate_btn = gr.Button("🎨 Generate Design", variant="primary", size="lg")
                
                # Prompt Preview Section
                with gr.Accordion("📋 Prompt Preview", open=False):
                    positive_prompt_preview = gr.Textbox(
                        label="✅ Positive Prompt",
                        lines=4,
                        interactive=False,
                        value="Select room type and style to see prompt preview"
                    )
                    negative_prompt_preview = gr.Textbox(
                        label="❌ Negative Prompt", 
                        lines=3,
                        interactive=False,
                        value="Select room type and style to see negative prompt preview"
                    )
            
            with gr.Column(scale=1):
                # Output
                output_image = gr.Image(
                    label="AI Design Result",
                    height=400
                )
                result_message = gr.Textbox(
                    label="Status",
                    interactive=False,
                    value="Ready to generate design"
                )
        
        # Update prompt preview when inputs change
        def update_prompt_preview(room_type, design_style, inpainting_mode):
            pos_prompt, neg_prompt = get_prompt_preview(room_type, design_style, inpainting_mode)
            return pos_prompt, neg_prompt
        
        # Event handlers for prompt preview updates
        for input_component in [room_type, design_style, inpainting_mode]:
            input_component.change(
                fn=update_prompt_preview,
                inputs=[room_type, design_style, inpainting_mode],
                outputs=[positive_prompt_preview, negative_prompt_preview]
            )
        
        # Initial prompt preview update
        demo.load(
            fn=update_prompt_preview,
            inputs=[room_type, design_style, inpainting_mode],
            outputs=[positive_prompt_preview, negative_prompt_preview]
        )
        
        # Main generation event handler
        generate_btn.click(
            fn=design_space,
            inputs=[input_image, room_type, design_style, inpainting_mode, num_steps, guidance_scale, strength],
            outputs=[output_image, result_message]
        )
        
    
    return demo

if __name__ == "__main__":
    # Load models on startup
    print("Starting Spacely AI Interior Designer...")
    
    # Create and launch interface
    demo = create_interface()
    demo.launch(share=True)