Update models.py

models.py

@@ -6,24 +6,29 @@ from diffusers import (
     PNDMScheduler,
     EulerDiscreteScheduler
 )
-from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
+from PIL import Image, ImageFilter, ImageEnhance
 import numpy as np
-from typing import List, Optional, Union, Tuple
-import os
+from typing import List, Union, Optional
+import cv2
 
 class InteriorDesignerPro:
     def __init__(self):
-        """Инициализация моделей для дизайна интерьеров"""
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model_name = "stabilityai/stable-diffusion-2-1"
         
         # Определяем мощность GPU
-        self.is_powerful_gpu = self._check_gpu_capability()
-        
-        print(f"🚀 Инициализация InteriorDesignerPro на {self.device}")
-        print(f"💪 Мощный GPU: {'Да' if self.is_powerful_gpu else 'Нет'}")
-        
-        # Основная модель для img2img
+        if self.device.type == "cuda":
+            gpu_name = torch.cuda.get_device_name(0).lower()
+            self.is_powerful_gpu = any(x in gpu_name for x in ["a100", "v100", "h100", "h200", "rtx 40", "rtx 30"])
+        else:
+            self.is_powerful_gpu = False
+            
+        print(f"🚀 Инициализация на {self.device}")
+        if self.device.type == "cuda":
+            print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
+            print(f"💪 Мощный GPU: {'Да' if self.is_powerful_gpu else 'Нет'}")
+            
+        # Основная модель
         self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
             self.model_name,
             torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
@@ -31,9 +36,8 @@ class InteriorDesignerPro:
             requires_safety_checker=False
         ).to(self.device)
         
-        # Модель для inpainting с обработкой ошибок
+        # Модель для inpainting (удаление объектов)
         try:
-            print("🎨 Загрузка модели для удаления объектов...")
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
                 torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
@@ -42,69 +46,64 @@ class InteriorDesignerPro:
                 local_files_only=False,
                 resume_download=True
             ).to(self.device)
-            print("✅ Модель inpainting загружена успешно")
+            print("✅ Inpainting модель загружена")
         except Exception as e:
-            print(f"⚠️ Не удалось загрузить модель inpainting: {e}")
-            print("📌 Используем основную модель как альтернативу")
-            # Используем основную модель для inpainting
-            self.inpaint_pipe = self.pipe
-        
+            print(f"⚠️ Не удалось загрузить inpainting модель: {e}")
+            print("Используем основную модель для inpainting")
+            self.inpaint_pipe = None
+            
         # Оптимизация памяти
         if self.device.type == "cuda":
             self.pipe.enable_attention_slicing()
             if hasattr(self.pipe, 'enable_xformers_memory_efficient_attention'):
                 try:
                     self.pipe.enable_xformers_memory_efficient_attention()
+                    if self.inpaint_pipe:
+                        self.inpaint_pipe.enable_xformers_memory_efficient_attention()
                     print("✅ xFormers оптимизация включена")
                 except:
-                    pass
+                    print("⚠️ xFormers недоступен")
                     
         # Настройка планировщиков
         self.schedulers = {
-            "DDIM": DDIMScheduler,
-            "PNDM": PNDMScheduler,
-            "Euler": EulerDiscreteScheduler
+            "ddim": DDIMScheduler.from_config(self.pipe.scheduler.config),
+            "pndm": PNDMScheduler.from_config(self.pipe.scheduler.config),
+            "euler": EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
         }
         
-        # Установка дефолтного планировщика
-        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        # Удаление объектов
+        self.object_remover = ObjectRemover(self.device)
         
-    def _check_gpu_capability(self):
-        """Проверка мощности GPU"""
-        if self.device.type != "cuda":
-            return False
-            
-        # Проверяем доступную память
-        try:
-            mem_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3
-            return mem_gb >= 12  # 12GB+ считаем мощным GPU
-        except:
-            return False
-    
-    def apply_style_pro(self, image: Image.Image, style_name: str, room_type: str, 
+    def apply_style_pro(self, image: Image.Image, style: str, room_type: str, 
                        strength: float = 0.75, quality: str = "balanced") -> Image.Image:
         """Применение стиля к изображению"""
         from design_styles import DESIGN_STYLES
         
-        if style_name not in DESIGN_STYLES:
-            raise ValueError(f"Неизвестный стиль: {style_name}")
+        if style not in DESIGN_STYLES:
+            style = "Современный минимализм"
+            
+        style_info = DESIGN_STYLES[style]
+        base_prompt = style_info.get("prompt", "")
+        negative_prompt = style_info.get("negative", "")
+        
+        # Добавляем специфику комнаты
+        if "room_specific" in style_info and room_type in style_info["room_specific"]:
+            room_prompt = style_info["room_specific"][room_type]
+            prompt = f"{room_prompt}, {base_prompt}, interior design, professional photo"
+        else:
+            prompt = f"{room_type} {base_prompt}, interior design, professional photo"
             
-        style = DESIGN_STYLES[style_name]
-        
         # Настройки качества
         quality_settings = {
-            "fast": {"steps": 20, "guidance": 7.5},
-            "balanced": {"steps": 35, "guidance": 8.5},
-            "ultra": {"steps": 50, "guidance": 10.0}
+            "fast": {"steps": 20, "guidance": 7.5, "scheduler": "euler"},
+            "balanced": {"steps": 35, "guidance": 8.5, "scheduler": "ddim"},
+            "ultra": {"steps": 50, "guidance": 10, "scheduler": "ddim"}
         }
         
         settings = quality_settings.get(quality, quality_settings["balanced"])
+        self.pipe.scheduler = self.schedulers[settings["scheduler"]]
         
-        # Генерация промпта
-        prompt = f"{room_type}, {style['prompt']}"
-        negative_prompt = style.get('negative', '')
-        
-        # Применение стиля
+        # Генерация
         result = self.pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -115,182 +114,197 @@ class InteriorDesignerPro:
         ).images[0]
         
         return result
-    
-    def remove_objects(self, image: Image.Image, mask: Image.Image, 
-                      inpaint_prompt: str = None) -> Image.Image:
-        """Удаление объектов с изображения"""
-        if inpaint_prompt is None:
-            inpaint_prompt = "empty room, clean walls, seamless texture"
-            
-        # Проверяем тип модели
-        if hasattr(self.inpaint_pipe, 'components') and 'vae' in self.inpaint_pipe.components:
-            # Это настоящая inpainting модель
-            result = self.inpaint_pipe(
-                prompt=inpaint_prompt,
-                negative_prompt="objects, furniture, people",
-                image=image,
-                mask_image=mask,
-                strength=0.99,
-                num_inference_steps=50,
-                guidance_scale=7.5
-            ).images[0]
-        else:
-            # Fallback на img2img с маской
-            # Создаем композитное изображение
-            masked_image = Image.composite(
-                Image.new('RGB', image.size, (255, 255, 255)),
-                image,
-                mask.convert('L').point(lambda x: 0 if x > 128 else 255)
-            )
-            
-            result = self.pipe(
-                prompt=inpaint_prompt,
-                negative_prompt="objects, furniture, people",
-                image=masked_image,
-                strength=0.95,
-                num_inference_steps=50,
-                guidance_scale=7.5
-            ).images[0]
-            
-        return result
-    
-    def create_variations(self, image: Image.Image, num_variations: int = 4) -> List[Image.Image]:
-        """Создание вариаций дизайна"""
+        
+    def create_variations(self, image: Image.Image, num_variations: int = 4,
+                         strength: float = 0.5) -> List[Image.Image]:
+        """Создание вариаций изображения"""
         variations = []
         
         prompts = [
-            "modern interior design, professional photo",
-            "cozy interior design, warm lighting",
-            "luxury interior design, high-end materials",
-            "minimalist interior design, clean space"
+            "modern interior design, bright and airy",
+            "cozy warm interior, soft lighting",
+            "minimalist clean space, neutral colors",
+            "luxury elegant interior, premium materials"
         ]
         
         for i in range(min(num_variations, len(prompts))):
-            variation = self.pipe(
+            var = self.pipe(
                 prompt=prompts[i],
                 image=image,
-                strength=0.5 + (i * 0.1),
-                num_inference_steps=30,
+                strength=strength,
+                num_inference_steps=25,
                 guidance_scale=7.5
             ).images[0]
-            variations.append(variation)
+            variations.append(var)
             
         return variations
-    
+        
     def create_hdr_lighting(self, image: Image.Image, intensity: float = 0.3) -> Image.Image:
         """Улучшение освещения (HDR эффект)"""
         # Конвертируем в numpy
-        img_array = np.array(image).astype(np.float32)
+        img_array = np.array(image)
         
         # Создаем HDR эффект
-        # Увеличиваем контраст в светлых областях
-        bright_mask = img_array > 200
-        img_array[bright_mask] = img_array[bright_mask] * (1 + intensity)
-        
-        # Усиливаем тени
-        dark_mask = img_array < 50
-        img_array[dark_mask] = img_array[dark_mask] * (1 - intensity * 0.5)
+        # Светлая версия
+        light = cv2.convertScaleAbs(img_array, alpha=1.2, beta=30)
+        # Темная версия  
+        dark = cv2.convertScaleAbs(img_array, alpha=0.8, beta=-30)
         
-        # Нормализация
-        img_array = np.clip(img_array, 0, 255).astype(np.uint8)
+        # Комбинируем
+        hdr = cv2.addWeighted(img_array, 1-intensity, light, intensity/2, 0)
+        hdr = cv2.addWeighted(hdr, 1, dark, intensity/2, 0)
         
-        # Применяем дополнительные фильтры
-        enhanced = Image.fromarray(img_array)
-        enhanced = enhanced.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
+        # Улучшаем контраст
+        lab = cv2.cvtColor(hdr, cv2.COLOR_RGB2LAB)
+        l, a, b = cv2.split(lab)
+        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
+        l = clahe.apply(l)
+        enhanced = cv2.merge([l, a, b])
+        result = cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB)
         
-        # Настройка яркости и контраста
-        enhancer = ImageEnhance.Contrast(enhanced)
-        enhanced = enhancer.enhance(1.1)
+        return Image.fromarray(result)
         
-        return enhanced
-    
     def enhance_details(self, image: Image.Image) -> Image.Image:
         """Улучшение деталей изображения"""
         # Увеличиваем резкость
-        enhanced = image.filter(ImageFilter.SHARPEN)
-        enhanced = enhanced.filter(ImageFilter.UnsharpMask(radius=2, percent=200, threshold=3))
+        enhancer = ImageEnhance.Sharpness(image)
+        sharp = enhancer.enhance(1.5)
         
-        # Улучшаем цвета
-        color_enhancer = ImageEnhance.Color(enhanced)
-        enhanced = color_enhancer.enhance(1.1)
+        # Улучшаем детали
+        detail_filter = ImageFilter.DETAIL
+        detailed = sharp.filter(detail_filter)
         
-        # Увеличиваем четкость
-        sharpness_enhancer = ImageEnhance.Sharpness(enhanced)
-        enhanced = sharpness_enhancer.enhance(1.2)
+        # Легкий unsharp mask
+        blurred = detailed.filter(ImageFilter.GaussianBlur(radius=1))
+        img_array = np.array(detailed)
+        blur_array = np.array(blurred) 
+        
+        # Применяем маску
+        sharpened = img_array + 0.5 * (img_array - blur_array)
+        sharpened = np.clip(sharpened, 0, 255).astype(np.uint8)
+        
+        return Image.fromarray(sharpened)
         
-        return enhanced
-    
     def change_element(self, image: Image.Image, element: str, value: str, 
-                      strength: float = 0.7) -> Image.Image:
-        """Изменение отдельного элемента интерьера"""
+                      strength: float = 0.5) -> Image.Image:
+        """Изменение конкретного элемента интерьера"""
         from design_styles import ROOM_ELEMENTS
         
         if element not in ROOM_ELEMENTS:
-            raise ValueError(f"Неизвестный элемент: {element}")
+            return image
             
         element_info = ROOM_ELEMENTS[element]
+        prompt_add = element_info.get("prompt_add", "")
         
-        # Создаем промпт для изменения
-        prompt = f"interior design, {element_info['prompt_add']}, {value}"
-        negative_prompt = "blurry, distorted"
+        prompt = f"{element} {value}, {prompt_add}, interior design"
+        negative = "bad quality, distorted"
         
-        # Применяем изменения
         result = self.pipe(
             prompt=prompt,
-            negative_prompt=negative_prompt,
+            negative_prompt=negative,
             image=image,
             strength=strength,
-            num_inference_steps=40,
+            num_inference_steps=30,
             guidance_scale=8.0
         ).images[0]
         
         return result
-    
+        
+    def remove_objects(self, image: Image.Image, mask: Image.Image, 
+                      inpaint_prompt: str = None) -> Image.Image:
+        """Удаление объектов с изображения"""
+        if self.inpaint_pipe is None:
+            # Fallback на обычную модель
+            return self.remove_objects_fallback(image, mask, inpaint_prompt)
+            
+        if inpaint_prompt is None:
+            inpaint_prompt = "empty room, clean space, seamless background"
+            
+        # Убеждаемся что маска в правильном формате
+        if mask.mode != "L":
+            mask = mask.convert("L")
+            
+        # Inpainting
+        result = self.inpaint_pipe(
+            prompt=inpaint_prompt,
+            negative_prompt="furniture, objects, people",
+            image=image,
+            mask_image=mask,
+            strength=0.99,
+            num_inference_steps=50,
+            guidance_scale=7.5
+        ).images[0]
+        
+        return result
+        
+    def remove_objects_fallback(self, image: Image.Image, mask: Image.Image,
+                               inpaint_prompt: str = None) -> Image.Image:
+        """Альтернативный метод удаления объектов"""
+        # Используем OpenCV inpainting
+        img_array = np.array(image)
+        mask_array = np.array(mask.convert("L"))
+        
+        # Расширяем маску для лучшего результата
+        kernel = np.ones((5,5), np.uint8)
+        mask_array = cv2.dilate(mask_array, kernel, iterations=2)
+        
+        # Inpainting
+        result = cv2.inpaint(img_array, mask_array, 3, cv2.INPAINT_TELEA)
+        
+        # Постобработка через img2img для улучшения
+        result_img = Image.fromarray(result)
+        
+        if inpaint_prompt is None:
+            inpaint_prompt = "clean empty space, seamless texture"
+            
+        enhanced = self.pipe(
+            prompt=inpaint_prompt,
+            image=result_img,
+            strength=0.3,
+            num_inference_steps=20,
+            guidance_scale=7.5
+        ).images[0]
+        
+        return enhanced
+
     def create_style_comparison(self, image: Image.Image, styles: List[str], 
                                room_type: str = "living room") -> Image.Image:
-        """Создание сравнения стилей"""
+        """Создание сравнения нескольких стилей"""
         styled_images = []
         
+        # Генерируем для каждого стиля
         for style in styles:
-            try:
-                styled = self.apply_style_pro(
-                    image, style, room_type, 
-                    strength=0.75, quality="fast"
-                )
-                styled_images.append((styled, style))
-            except Exception as e:
-                print(f"Ошибка при применении стиля {style}: {e}")
-                
-        if not styled_images:
-            return image
+            styled = self.apply_style_pro(
+                image, style, room_type, 
+                strength=0.75, 
+                quality="fast"  # Быстрый режим для множественной генерации
+            )
+            styled_images.append(styled)
             
-        # Используем метод из процессора
-        if hasattr(self, '_create_comparison_grid'):
-            return self._create_comparison_grid(styled_images)
-        else:
-            # Простая сетка
-            return self._simple_grid(styled_images)
-    
-    def _simple_grid(self, images_with_titles: List[Tuple[Image.Image, str]]) -> Image.Image:
-        """Создание простой сетки изображений"""
-        if not images_with_titles:
+        # Создаем сетку
+        return self._create_comparison_grid(styled_images, styles)
+        
+    def _create_comparison_grid(self, images: List[Image.Image], 
+                               titles: List[str]) -> Image.Image:
+        """Создание сетки из изображений"""
+        if not images:
             return None
             
-        images = [img for img, _ in images_with_titles]
-        titles = [title for _, title in images_with_titles]
+        # Определяем размер сетки
+        n = len(images)
+        cols = min(3, n)
+        rows = (n + cols - 1) // cols
         
-        # Размеры
+        # Размер одного изображения
         img_width, img_height = images[0].size
-        cols = min(3, len(images))
-        rows = (len(images) + cols - 1) // cols
         
-        # Создаем сетку
-        grid_width = img_width * cols
-        grid_height = img_height * rows
+        # Создаем холст
+        grid_width = cols * img_width
+        grid_height = rows * img_height
         grid = Image.new('RGB', (grid_width, grid_height), 'white')
         
         # Размещаем изображения
-        for idx, (img, title) in enumerate(images_with_titles):
+        for idx, (img, title) in enumerate(zip(images, titles)):
             row = idx // cols
             col = idx % cols
             x = col * img_width
@@ -300,25 +314,32 @@ class InteriorDesignerPro:
         return grid
 
 
-# Дополнительный класс для удаления объектов
 class ObjectRemover:
+    """Класс для удаления объектов"""
     def __init__(self, device):
         self.device = device
         
-    def generate_mask_from_text(self, image: Image.Image, text: str, precision: float) -> Image.Image:
+    def generate_mask_from_text(self, image: Image.Image, text: str, 
+                               precision: float = 0.3) -> Image.Image:
         """Генерация маски на основе текстового описания"""
         # Простая реализация - создаем маску в центре
         width, height = image.size
         mask = Image.new('L', (width, height), 0)
+        
+        # Создаем примерную область в центре
+        from PIL import ImageDraw
         draw = ImageDraw.Draw(mask)
         
         # Размер области зависит от precision
-        margin = int(min(width, height) * (1 - precision) / 2)
+        margin_x = int(width * (0.5 - precision))
+        margin_y = int(height * (0.5 - precision))
         
-        # Рисуем белую область в центре
-        draw.ellipse([margin, margin, width-margin, height-margin], fill=255)
+        draw.ellipse(
+            [margin_x, margin_y, width - margin_x, height - margin_y],
+            fill=255
+        )
         
-        # Размываем для плавных краев
+        # Размываем маску
         mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
         
         return mask