app.py

"""
Interior AI Designer API
Hugging Face Spaces deployment with public API endpoint
"""

import os
import random
import time
import gradio as gr
import numpy as np
import spaces  # Required for Hugging Face Spaces GPU
import torch
import gc
from PIL import Image
from diffusers import (
    ControlNetModel,
    DPMSolverMultistepScheduler,
    StableDiffusionControlNetPipeline,
)
from controlnet_aux_local import NormalBaeDetector

MAX_SEED = np.iinfo(np.int32).max

# ============================================================
# Model Loading (runs once at startup)
# ============================================================

class Preprocessor:
    MODEL_ID = "lllyasviel/Annotators"

    def __init__(self):
        self.model = None
        self.name = ""

    def load(self, name: str) -> None:
        if name == self.name:
            return
        elif name == "NormalBae":
            print("Loading NormalBae")
            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
            torch.cuda.empty_cache()
            self.name = name
        else:
            raise ValueError(f"Unknown preprocessor: {name}")

    def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
        return self.model(image, **kwargs)


# Load models at startup
if gr.NO_RELOAD:
    print("CUDA version:", torch.version.cuda)
    print("Loading models...")

    # ControlNet
    model_id = "lllyasviel/control_v11p_sd15_normalbae"
    controlnet = ControlNetModel.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
    ).to("cuda")

    # Scheduler
    scheduler = DPMSolverMultistepScheduler.from_pretrained(
        "ashllay/stable-diffusion-v1-5-archive",
        solver_order=2,
        subfolder="scheduler",
        use_karras_sigmas=True,
        final_sigmas_type="sigma_min",
        algorithm_type="sde-dpmsolver++",
        prediction_type="epsilon",
        thresholding=False,
        denoise_final=True,
        torch_dtype=torch.float16,
    )

    # Stable Diffusion Pipeline
    base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
    
    pipe = StableDiffusionControlNetPipeline.from_single_file(
        base_model_url,
        safety_checker=None,
        controlnet=controlnet,
        scheduler=scheduler,
        torch_dtype=torch.float16,
    ).to("cuda")

    # Preprocessor
    preprocessor = Preprocessor()
    preprocessor.load("NormalBae")

    # Optional: Load textual inversions for better negative prompts
    try:
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
        pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
    except Exception as e:
        print(f"Could not load textual inversions: {e}")

    pipe.to("cuda")
    torch.cuda.empty_cache()
    gc.collect()
    print(f"Models loaded! CUDA memory: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")


# ============================================================
# Style Definitions
# ============================================================

STYLE_LIST = [
    {"name": "None", "prompt": ""},
    {"name": "Minimalistic", "prompt": "Minimalist interior design,clean lines,neutral colors,uncluttered space,functional furniture,lots of natural light"},
    {"name": "Boho", "prompt": "Bohemian chic interior,eclectic mix of patterns and textures,vintage furniture,plants,woven textiles,warm earthy colors"},
    {"name": "Farmhouse", "prompt": "Modern farmhouse interior,rustic wood elements,shiplap walls,neutral color palette,industrial accents,cozy textiles"},
    {"name": "Saudi Prince", "prompt": "Opulent gold interior,luxurious ornate furniture,crystal chandeliers,rich fabrics,marble floors,intricate Arabic patterns"},
    {"name": "Neoclassical", "prompt": "Neoclassical interior design,elegant columns,ornate moldings,symmetrical layout,refined furniture,muted color palette"},
    {"name": "Eclectic", "prompt": "Eclectic interior design,mix of styles and eras,bold color combinations,diverse furniture pieces,unique art objects"},
    {"name": "Parisian", "prompt": "Parisian apartment interior,all-white color scheme,ornate moldings,herringbone wood floors,elegant furniture,large windows"},
    {"name": "Hollywood", "prompt": "Hollywood Regency interior,glamorous and luxurious,bold colors,mirrored surfaces,velvet upholstery,gold accents"},
    {"name": "Scandinavian", "prompt": "Scandinavian interior design,light wood tones,white walls,minimalist furniture,cozy textiles,hygge atmosphere"},
    {"name": "Beach", "prompt": "Coastal beach house interior,light blue and white color scheme,weathered wood,nautical accents,sheer curtains,ocean view"},
    {"name": "Japanese", "prompt": "Traditional Japanese interior,tatami mats,shoji screens,low furniture,zen garden view,minimalist decor,natural materials"},
    {"name": "Midcentury Modern", "prompt": "Mid-century modern interior,1950s-60s style furniture,organic shapes,warm wood tones,bold accent colors,large windows"},
    {"name": "Retro Futurism", "prompt": "Neon (atompunk world) retro cyberpunk background"},
    {"name": "Texan", "prompt": "Western cowboy interior,rustic wood beams,leather furniture,cowhide rugs,antler chandeliers,southwestern patterns"},
    {"name": "Matrix", "prompt": "Futuristic cyberpunk interior,neon accent lighting,holographic plants,sleek black surfaces,advanced gaming setup,transparent screens,Blade Runner inspired decor,high-tech minimalist furniture"},
]

STYLES = {s["name"]: s["prompt"] for s in STYLE_LIST}
STYLE_NAMES = list(STYLES.keys())


# ============================================================
# Core Processing Function (API Endpoint)
# ============================================================

@spaces.GPU(duration=20)
@torch.inference_mode()
def redesign_interior(
    image: Image.Image,
    style: str = "Minimalistic",
    custom_prompt: str = "",
    num_steps: int = 15,
    guidance_scale: float = 5.5,
    seed: int = -1,
    image_resolution: int = 768,
) -> Image.Image:
    """
    Redesign an interior image with the specified style.
    
    Args:
        image: Input room/interior image (PIL Image)
        style: Design style name (e.g., "Minimalistic", "Boho", "Japanese")
        custom_prompt: Additional custom prompt to append
        num_steps: Number of inference steps (default: 15)
        guidance_scale: Guidance scale for generation (default: 5.5)
        seed: Random seed (-1 for random)
        image_resolution: Output image resolution (default: 768)
    
    Returns:
        Redesigned interior image (PIL Image)
    """
    # Set seed
    if seed == -1:
        seed = random.randint(0, MAX_SEED)
    generator = torch.cuda.manual_seed(seed)
    
    # Preprocess image with NormalBae
    preprocessor.load("NormalBae")
    control_image = preprocessor(
        image=image,
        image_resolution=image_resolution,
        detect_resolution=image_resolution,
    )
    
    # Build prompt
    base_prompt = "Photo from Pinterest of"
    style_prompt = STYLES.get(style, "")
    additional_prompt = "design-style interior designed (interior space), tungsten white balance, captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length"
    
    if style_prompt:
        prompt = f"{base_prompt} {style_prompt} {custom_prompt}, {additional_prompt}"
    else:
        prompt = f"{base_prompt} {custom_prompt}, {additional_prompt}" if custom_prompt else f"boho chic interior, {additional_prompt}"
    
    negative_prompt = "EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)"
    
    print(f"Prompt: {prompt}")
    print(f"Style: {style}, Seed: {seed}")
    
    # Generate
    start = time.time()
    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=guidance_scale,
        num_images_per_prompt=1,
        num_inference_steps=num_steps,
        generator=generator,
        image=control_image,
    ).images[0]
    
    print(f"Generation completed in {time.time() - start:.2f}s")
    torch.cuda.empty_cache()
    
    return result


# ============================================================
# Gradio Interface
# ============================================================

with gr.Blocks() as demo:
    gr.Markdown("# 🏠 Interior AI Designer")
    gr.Markdown("Upload a room photo and select a design style to reimagine your space!")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="Upload Room Image",
                type="pil",
                sources=["upload", "clipboard"],
            )
            style_dropdown = gr.Dropdown(
                label="Design Style",
                choices=STYLE_NAMES,
                value="Minimalistic",
            )
            custom_prompt = gr.Textbox(
                label="Custom Prompt (optional)",
                placeholder="Add specific details like 'with plants' or 'blue accents'",
            )
            
            with gr.Accordion("Advanced Options", open=False):
                num_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10, maximum=50, value=15, step=1,
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0, maximum=20.0, value=5.5, step=0.5,
                )
                seed = gr.Slider(
                    label="Seed (-1 for random)",
                    minimum=-1, maximum=MAX_SEED, value=-1, step=1,
                )
                image_resolution = gr.Slider(
                    label="Resolution",
                    minimum=512, maximum=1024, value=768, step=128,
                )
            
            generate_btn = gr.Button("🎨 Redesign Interior", variant="primary", size="lg")
        
        with gr.Column():
            output_image = gr.Image(label="Redesigned Interior", type="pil")
    
    # Examples
    gr.Examples(
        examples=[
            ["Minimalistic"],
            ["Boho"],
            ["Japanese"],
            ["Scandinavian"],
            ["Matrix"],
        ],
        inputs=[style_dropdown],
        label="Try these styles",
    )
    
    # Connect the button to the function
    generate_btn.click(
        fn=redesign_interior,
        inputs=[input_image, style_dropdown, custom_prompt, num_steps, guidance_scale, seed, image_resolution],
        outputs=output_image,
        api_name="redesign",  # This enables the API endpoint
    )
    
    gr.Markdown("""
    ---
    ### 📡 API Usage
    
    This app exposes a public API! You can call it programmatically:
    
    ```python
    from gradio_client import Client
    
    client = Client("YOUR_USERNAME/interior-ai-designer")
    result = client.predict(
        image="path/to/room.jpg",
        style="Minimalistic",
        custom_prompt="",
        num_steps=15,
        guidance_scale=5.5,
        seed=-1,
        image_resolution=768,
        api_name="/redesign"
    )
    print(result)  # Path to output image
    ```
    """)


# Launch
if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        server_name="0.0.0.0",
        server_port=7860,
    )