app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageEnhance
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import spaces
import os

# --- 1. БАЗЫ ЗНАНИЙ ---
ROOM_PROMPTS = {
    "Гостиная": "living room interior, sofa area, coffee table, lounge",
    "Кухня": "kitchen interior, kitchen cabinets, countertop, sink, appliances",
    "Спальня": "bedroom interior, bed with pillows, headboard, cozy",
    "Ванная / Санузел": "bathroom interior, toilet, sink, mirror, tiled walls, faucet",
    "Детская": "kids room interior, playroom, toys, single bed, colorful",
    "Кабинет / Офис": "home office interior, desk, office chair, computer, shelves",
    "Прихожая / Коридор": "hallway interior, corridor, entrance, wardrobe, mirror",
    "Общественное (Кафе/Лобби)": "public space interior, cafe interior, tables and chairs, restaurant"
}

STYLE_PROMPTS = {
    "Scandi (Скандинавский)": "scandinavian style, white walls, light oak wood, beige palette, soft sunlight, hygge, natural materials",
    "Loft (Индустриальный)": "industrial loft style, concrete walls, exposed brick, black metal accents, leather textures, dramatic lighting, cinematic",
    "Neoclassic (Неоклассика)": "modern classic interior, luxury, wall moldings, velvet furniture, marble floors, gold details, crystal chandelier, elegant",
    "Minimalism (Минимализм)": "minimalism style, ultra modern, clean lines, monolithic forms, hidden lighting, decluttered, monochromatic grey and white",
    "Japandi (Джапанди)": "japandi style, wabi-sabi, warm earth tones, raw wood, textured plaster walls, organic shapes, zen atmosphere, soft shadows",
    "Dark Luxury (Муди)": "dark moody interior, walnut wood, dark grey walls, warm dim lighting, fireplace atmosphere, rich textures, expensive"
}

# --- 2. ЗАГРУЗКА МОДЕЛИ ---
pipe = None

def load_pipeline():
    global pipe
    if pipe is None:
        print("Загрузка модели...")
        controlnet = ControlNetModel.from_pretrained(
            "lllyasviel/sd-controlnet-canny", 
            torch_dtype=torch.float32
        )
        pipe = StableDiffusionControlNetPipeline.from_pretrained(
            "SG161222/Realistic_Vision_V5.1_noVAE", 
            controlnet=controlnet, 
            torch_dtype=torch.float32
        )
        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
        try:
            pipe.to("cuda")
        except:
            pass 
    return pipe

load_pipeline()

# --- 3. ОБРАБОТКА ---
def boost_contrast(image, factor=2.0):
    enhancer = ImageEnhance.Contrast(image)
    return enhancer.enhance(factor)

def process_canny(image):
    if image is None: return None
    contrasted_image = boost_contrast(image, factor=2.0)
    image_np = np.array(contrasted_image)
    image_canny = cv2.Canny(image_np, 50, 150)
    image_canny = image_canny[:, :, None]
    image_canny = np.concatenate([image_canny, image_canny, image_canny], axis=2)
    return Image.fromarray(image_canny)

# --- 4. ГЕНЕРАЦИЯ ---
@spaces.GPU(duration=60)
def generate(image_path, room_type, style_name, strength, seed):
    if image_path is None:
        return None, None, None
    
    image = Image.open(image_path).convert("RGB")
    original_filename = os.path.splitext(os.path.basename(image_path))[0]
    
    if max(image.size) > 512:
        ratio = 512 / max(image.size)
        new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio))
        image = image.resize(new_size, Image.LANCZOS)
    
    selected_room = ROOM_PROMPTS[room_type]
    selected_style = STYLE_PROMPTS[style_name]
    
    final_prompt = f"{selected_room}, {selected_style}, architectural photography, masterpiece, 8k"
    negative_prompt = "blurry, low quality, distorted perspective, bad anatomy, worst quality, watermark, signature"

    canny_image = process_canny(image)
    
    generator = torch.manual_seed(int(seed))
    result = pipe(
        prompt=final_prompt,
        negative_prompt=negative_prompt,
        image=canny_image,
        num_inference_steps=20,
        generator=generator,
        controlnet_conditioning_scale=float(strength)
    ).images[0]
    
    output_filename = f"3D_CROSS.{original_filename}.jpg"
    result.save(output_filename, format="JPEG", quality=95)
    
    return result, canny_image, output_filename

# --- 5. JS СКРИПТЫ ---
# Шеринг
share_js = """
async () => {
    const imgElement = document.querySelector('.result-image img');
    if (!imgElement) {
        alert('Сначала сгенерируйте изображение!');
        return;
    }
    try {
        const canvas = document.createElement('canvas');
        canvas.width = imgElement.naturalWidth;
        canvas.height = imgElement.naturalHeight;
        const ctx = canvas.getContext('2d');
        ctx.drawImage(imgElement, 0, 0);

        canvas.toBlob(async (blob) => {
            if (!blob) {
                alert('Ошибка обработки изображения.');
                return;
            }
            const file = new File([blob], '3D_CROSS_render.jpg', { type: 'image/jpeg' });
            if (navigator.share) {
                try {
                    await navigator.share({
                        title: '3D CROSS Render',
                        text: 'Смотри, какой дизайн я сделал в 3D CROSS!',
                        files: [file]
                    });
                } catch (err) {
                    console.log('Отмена.');
                }
            } else {
                try {
                    const clipboardItem = new ClipboardItem({'image/png': blob});
                    await navigator.clipboard.write([clipboardItem]);
                    alert('Изображение скопировано в буфер! (Ctrl+V)');
                } catch (err) {
                    alert('Ваш браузер не поддерживает отправку. Скачайте файл вручную.');
                }
            }
        }, 'image/jpeg', 0.95);
    } catch (error) {
        console.error('Ошибка JS:', error);
        alert('Ошибка. Попробуйте скачать файл.');
    }
}
"""

# --- 6. ИНТЕРФЕЙС ---
css = """
/* Скрываем кнопки Download и Share */
.gradio-container button[aria-label="Download"],
.gradio-container button[aria-label="Share"] {
    display: none !important;
}

/* Картинка */
.result-image img {
    width: 100% !important;
    height: auto !important;
    max-height: 600px !important; 
    object-fit: contain !important;
}

/* Fullscreen fix */
.gradio-image-preview {
    background-color: rgba(10, 10, 10, 0.98) !important;
    z-index: 9999 !important;
}
.gradio-image-preview img {
    height: 95vh !important;
    width: auto !important;
    max-width: 95vw !important;
    object-fit: contain !important;
    margin: auto !important;
    display: block !important;
}

/* Убираем отступы вокруг приложения, чтобы оно красиво встало в iframe */
.container { max-width: 100% !important; margin: 0 !important; padding: 10px !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo:
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Исходник", type="filepath", height=400)
            
            with gr.Accordion("Настройки проекта", open=True):
                room_dropdown = gr.Dropdown(choices=list(ROOM_PROMPTS.keys()), value="Гостиная", label="📂 Тип помещения")
                style_dropdown = gr.Dropdown(choices=list(STYLE_PROMPTS.keys()), value="Scandi (Скандинавский)", label="🎨 Стиль")
                strength_slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="Жесткость")
                seed_number = gr.Number(value=42, label="Seed", precision=0)
            
            run_btn = gr.Button("СГЕНЕРИРОВАТЬ 🚀", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            result_image = gr.Image(
                label="Результат (Кликни для увеличения)", 
                type="pil", 
                height=600,
                elem_classes="result-image"
            )
            
            with gr.Row():
                download_btn = gr.DownloadButton("📥 СКАЧАТЬ JPG", variant="secondary")
                share_btn = gr.Button("🔗 ОТПРАВИТЬ / КОПИРОВАТЬ") 
            
            with gr.Accordion("Debug", open=False):
                canny_debug = gr.Image(label="Карта линий", type="pil")

    # Генерация
    run_btn.click(
        fn=generate,
        inputs=[input_image, room_dropdown, style_dropdown, strength_slider, seed_number],
        outputs=[result_image, canny_debug, download_btn]
    )
    
    # Кнопки JS
    share_btn.click(None, [], [], js=share_js)

if __name__ == "__main__":
    demo.launch()