models.py

import torch
from diffusers import (
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    DDIMScheduler,
    PNDMScheduler,
    EulerDiscreteScheduler,
    DPMSolverMultistepScheduler
)
from PIL import Image, ImageFilter, ImageEnhance
import numpy as np
import cv2

class InteriorDesignerPro:
    def __init__(self):
        self.device = torch.device("cuda")  # ТОЛЬКО GPU!
        self.model_name = "RealVisXL V4.0"

        # Проверка GPU
        gpu_name = torch.cuda.get_device_name(0)
        self.is_powerful_gpu = any(gpu in gpu_name for gpu in ['A100', 'H100', 'RTX 4090', 'RTX 3090', 'T4', 'A10G'])

        # Основная модель - RealVisXL V4 
        print(f"Loading {self.model_name} on {gpu_name}...")
        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
            "SG161222/RealVisXL_V4.0",
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16"
        ).to(self.device)

        # БЕЗ enable_model_cpu_offload() и enable_vae_slicing() - они замедляют H200!

        # Настройка scheduler для качества
        self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)

        # Inpainting модель
        try:
            self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                "stabilityai/stable-diffusion-2-inpainting",
                torch_dtype=torch.float16,
                safety_checker=None,
                requires_safety_checker=False,
                local_files_only=False,
                resume_download=True
            ).to(self.device)
            print("Inpainting model loaded")
        except Exception as e:
            print(f"Warning: Could not load inpainting model: {e}")
            print("Using img2img as fallback for object removal")
            self.inpaint_pipe = None

    @torch.inference_mode()
    def apply_style_pro(self, image, style_name, room_type, strength=0.75, quality="balanced", custom_prompt=None, custom_negative=None):
        """Применение стиля к изображению"""
        from design_styles import DESIGN_STYLES

        # Ресайз для скорости
        original_size = image.size
        if image.width > 768 or image.height > 768:
            image.thumbnail((768, 768), Image.Resampling.LANCZOS)

        if style_name == "custom" and custom_prompt:
            # Кастомный промпт
            full_prompt = custom_prompt
            negative = custom_negative or "low quality, blurry"
        else:
            # Предустановленный стиль
            style = DESIGN_STYLES.get(style_name, DESIGN_STYLES["Современный минимализм"])
            room_specific = style.get("room_specific", {}).get(room_type, "")
            full_prompt = f"{style['prompt']}, {room_specific}, {room_type} interior design, professional photo, high quality, 8k, photorealistic"
            negative = style.get("negative", "low quality, blurry")

        # Настройки качества - оптимизированные для H200
        quality_settings = {
            "fast": {"steps": 15, "guidance": 6.0},
            "balanced": {"steps": 20, "guidance": 7.0},
            "ultra": {"steps": 30, "guidance": 8.0}
        }

        settings = quality_settings.get(quality, quality_settings["balanced"])

        # Генерация с SDXL
        result = self.pipe(
            prompt=full_prompt,
            prompt_2=full_prompt,  # Для SDXL
            negative_prompt=negative,
            negative_prompt_2=negative,  # Для SDXL
            image=image,
            strength=strength,
            num_inference_steps=settings["steps"],
            guidance_scale=settings["guidance"],
            # SDXL параметры - оптимизированные
            original_size=(768, 768),
            target_size=(768, 768)
        ).images[0]

        # Возвращаем к оригинальному размеру если нужно
        if result.size != original_size and max(original_size) <= 1024:
            result = result.resize(original_size, Image.Resampling.LANCZOS)

        return result

    def create_variations(self, image, num_variations=4):
        """Создание вариаций дизайна"""
        variations = []
        base_seed = torch.randint(0, 1000000, (1,)).item()

        # Ресайз для скорости
        if image.width > 768 or image.height > 768:
            image.thumbnail((768, 768), Image.Resampling.LANCZOS)

        for i in range(num_variations):
            torch.manual_seed(base_seed + i)

            var = self.pipe(
                prompt="interior design variation, same style, different details",
                prompt_2="interior design variation, same style, different details",
                image=image,
                strength=0.4 + (i * 0.05),
                num_inference_steps=20,  # Меньше шагов для скорости
                guidance_scale=6.0
            ).images[0]

            variations.append(var)

        return variations

    def create_hdr_lighting(self, image, intensity=0.3):
        """Улучшение освещения в стиле HDR"""
        # Конвертируем в numpy
        img_array = np.array(image)

        # Применяем CLAHE для улучшения контраста
        lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)

        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l_clahe = clahe.apply(l)

        enhanced_lab = cv2.merge([l_clahe, a, b])
        enhanced_rgb = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)

        # Смешиваем с оригиналом
        result = cv2.addWeighted(img_array, 1-intensity, enhanced_rgb, intensity, 0)

        return Image.fromarray(result)

    def enhance_details(self, image):
        """Улучшение деталей изображения"""
        # Увеличиваем резкость
        enhancer = ImageEnhance.Sharpness(image)
        sharp = enhancer.enhance(1.5)

        # Немного увеличиваем контраст
        enhancer = ImageEnhance.Contrast(sharp)
        contrast = enhancer.enhance(1.1)

        return contrast

    def change_element(self, image, element, value, strength=0.7):
        """Изменение отдельного элемента интерьера"""
        from design_styles import ROOM_ELEMENTS

        # Ресайз для скорости
        original_size = image.size
        if image.width > 768 or image.height > 768:
            image.thumbnail((768, 768), Image.Resampling.LANCZOS)

        element_info = ROOM_ELEMENTS.get(element, {})
        prompt_add = element_info.get("prompt_add", element.lower())

        prompt = f"interior with {value} {prompt_add}, professional photo"
        negative = f"old {element}, damaged, ugly"

        result = self.pipe(
            prompt=prompt,
            prompt_2=prompt,  # Для SDXL
            negative_prompt=negative,
            negative_prompt_2=negative,
            image=image,
            strength=min(strength, 0.8),  # Ограничиваем для скорости
            num_inference_steps=20,  # Оптимизировано для H200
            guidance_scale=6.0
        ).images[0]

        # Возвращаем к оригинальному размеру
        if result.size != original_size:
            result = result.resize(original_size, Image.Resampling.LANCZOS)

        return result

    def create_style_comparison(self, image, styles, quality="fast"):
        """Создание сравнения стилей"""
        results = []

        # Настройки для быстрой генерации
        steps = 15 if quality == "fast" else 20

        for style in styles:
            styled = self.apply_style_pro(
                image, 
                style, 
                "living room",  # default
                strength=0.75,
                quality=quality
            )
            results.append((style, styled))

        return results


class ObjectRemover:
    """Класс для удаления объектов - оптимизированный"""

    def __init__(self, inpaint_pipe):
        self.pipe = inpaint_pipe
        self.device = torch.device("cuda")

    def remove_objects(self, image, mask):
        """Удаление объектов с изображения"""
        if self.pipe is None:
            # Fallback на простое заполнение
            return self.simple_inpaint(image, mask)

        try:
            # Используем inpainting pipeline с оптимизированными параметрами
            result = self.pipe(
                prompt="empty room interior, clean wall, seamless texture",
                negative_prompt="furniture, objects, people, clutter",
                image=image,
                mask_image=mask,
                strength=0.95,  # Немного меньше для скорости
                num_inference_steps=25,  # Оптимизировано!
                guidance_scale=5.0  # Меньше для скорости
            ).images[0]

            return result
        except Exception as e:
            print(f"Inpainting failed: {e}, using OpenCV fallback")
            return self.simple_inpaint(image, mask)

    def simple_inpaint(self, image, mask):
        """Простое заполнение через OpenCV - очень быстро"""
        img_array = np.array(image)
        mask_array = np.array(mask.convert('L'))

        # Инпейнтинг через OpenCV
        result = cv2.inpaint(img_array, mask_array, 3, cv2.INPAINT_TELEA)

        return Image.fromarray(result)

    def generate_mask_from_text(self, image, text_description, precision=0.3):
        """Генерация маски на основе текстового описания"""
        # Простая маска в центре (заглушка)
        # В реальности тут должен быть CLIP или SAM
        width, height = image.size
        mask = Image.new('L', (width, height), 0)

        # Создаем маску в центре
        center_x, center_y = width // 2, height // 2
        radius = int(min(width, height) * precision)

        # Рисуем круг
        from PIL import ImageDraw
        draw = ImageDraw.Draw(mask)
        draw.ellipse([center_x - radius, center_y - radius, 
                      center_x + radius, center_y + radius], fill=255)

        return mask