Create models.py

models.py

@@ -0,0 +1,270 @@
+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