Update models.py

models.py

@@ -2,128 +2,109 @@ import torch
 from diffusers import (
     StableDiffusionImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
-    DPMSolverMultistepScheduler,
-    EulerDiscreteScheduler,
     DDIMScheduler,
-    PNDMScheduler
+    PNDMScheduler,
+    EulerDiscreteScheduler
 )
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFilter, ImageEnhance
 import numpy as np
-from typing import List, Optional, Union
-import cv2
+from typing import List, Optional, Union, Tuple
+import os
 
 class InteriorDesignerPro:
-    """Профессиональный AI дизайнер интерьеров"""
-    
     def __init__(self):
-        """Инициализация моделей и параметров"""
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.is_gpu = torch.cuda.is_available()
-        self.is_powerful_gpu = False
-        
-        # Проверяем мощность GPU
-        if self.is_gpu:
-            gpu_name = torch.cuda.get_device_name(0).lower()
-            # H200, H100, A100, RTX 4090 и т.д.
-            powerful_gpus = ['h200', 'h100', 'a100', '4090', 'a6000']
-            self.is_powerful_gpu = any(gpu in gpu_name for gpu in powerful_gpus)
-        
-        # Базовая модель
-        self.model_id = "runwayml/stable-diffusion-v1-5"
-        self.model_name = "Stable Diffusion v1.5"
-        
-        # Инициализация пайплайнов
-        self.pipe = None
-        self.inpaint_pipe = None
-        self.safety_checker = None
-        
-        self._load_models()
-        
-    def _load_models(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
+        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            self.model_name,
+            torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(self.device)
+        
+        # Модель для inpainting с обработкой ошибок
         try:
-            # Основной пайплайн для img2img
-            self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-                self.model_id,
-                torch_dtype=torch.float16 if self.is_gpu else torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            self.pipe = self.pipe.to(self.device)
-            
-            # Пайплайн для inpainting
+            print("🎨 Загрузка модели для удаления объектов...")
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
-                torch_dtype=torch.float16 if self.is_gpu else torch.float32,
+                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
                 safety_checker=None,
-                requires_safety_checker=False
-            )
-            self.inpaint_pipe = self.inpaint_pipe.to(self.device)
-            
-            # Настройка планировщиков
-            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-                self.pipe.scheduler.config
-            )
-            
-            # Оптимизация памяти
-            if self.is_gpu:
-                self.pipe.enable_attention_slicing()
-                self.inpaint_pipe.enable_attention_slicing()
-                
-                # xFormers для ускорения (если доступно)
+                requires_safety_checker=False,
+                local_files_only=False,
+                resume_download=True
+            ).to(self.device)
+            print("✅ Модель inpainting загружена успешно")
+        except Exception as e:
+            print(f"⚠️ Не удалось загрузить модель inpainting: {e}")
+            print("📌 Используем основную модель как альтернативу")
+            # Используем основную модель для inpainting
+            self.inpaint_pipe = self.pipe
+        
+        # Оптимизация памяти
+        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()
-                    self.inpaint_pipe.enable_xformers_memory_efficient_attention()
+                    print("✅ xFormers оптимизация включена")
                 except:
                     pass
+                    
+        # Настройка планировщиков
+        self.schedulers = {
+            "DDIM": DDIMScheduler,
+            "PNDM": PNDMScheduler,
+            "Euler": EulerDiscreteScheduler
+        }
+        
+        # Установка дефолтного планировщика
+        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        
+    def _check_gpu_capability(self):
+        """Проверка мощности GPU"""
+        if self.device.type != "cuda":
+            return False
             
-        except Exception as e:
-            print(f"Ошибка загрузки моделей: {e}")
-            raise
+        # Проверяем доступную память
+        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(self, image: Image.Image, style: str, room_type: str = "living room",
-                   strength: float = 0.8) -> Image.Image:
+    def apply_style_pro(self, image: Image.Image, style_name: str, room_type: str, 
+                       strength: float = 0.75, quality: str = "balanced") -> Image.Image:
         """Применение стиля к изображению"""
         from design_styles import DESIGN_STYLES
         
-        style_data = DESIGN_STYLES.get(style, DESIGN_STYLES["Современный минимализм"])
-        
-        # Базовый промпт
-        base_prompt = f"{room_type}, {style_data['prompt']}, interior design photography, high quality, 8k"
-        negative_prompt = f"{style_data.get('negative', '')}, low quality, blurry, distorted"
-        
-        # Применяем стиль
-        result = self.pipe(
-            prompt=base_prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            strength=strength * style_data.get('strength', 0.75),
-            num_inference_steps=30,
-            guidance_scale=7.5
-        ).images[0]
-        
-        return result
-    
-    def apply_style_pro(self, image: Image.Image, style: str, room_type: str = "living room",
-                       strength: float = 0.8, quality: str = "balanced") -> Image.Image:
-        """Применение стиля с учетом качества"""
-        from design_styles import DESIGN_STYLES, get_detailed_prompt
-        
-        # Получаем детальный промпт
-        prompt = get_detailed_prompt(style, room_type, quality=quality)
-        style_data = DESIGN_STYLES.get(style, {})
-        negative_prompt = style_data.get('negative', 'low quality, blurry')
+        if style_name not in DESIGN_STYLES:
+            raise ValueError(f"Неизвестный стиль: {style_name}")
+            
+        style = DESIGN_STYLES[style_name]
         
-        # Настройки в зависимости от качества
+        # Настройки качества
         quality_settings = {
             "fast": {"steps": 20, "guidance": 7.5},
             "balanced": {"steps": 35, "guidance": 8.5},
-            "ultra": {"steps": 50, "guidance": 10}
+            "ultra": {"steps": 50, "guidance": 10.0}
         }
         
         settings = quality_settings.get(quality, quality_settings["balanced"])
         
-        # Применяем стиль
+        # Генерация промпта
+        prompt = f"{room_type}, {style['prompt']}"
+        negative_prompt = style.get('negative', '')
+        
+        # Применение стиля
         result = self.pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -135,109 +116,131 @@ class InteriorDesignerPro:
         
         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]:
         """Создание вариаций дизайна"""
         variations = []
         
-        # Базовые параметры для вариаций
-        strengths = [0.6, 0.7, 0.75, 0.8, 0.85, 0.9, 0.65, 0.82]
-        seeds = [42, 123, 456, 789, 1011, 1213, 1415, 1617]
+        prompts = [
+            "modern interior design, professional photo",
+            "cozy interior design, warm lighting",
+            "luxury interior design, high-end materials",
+            "minimalist interior design, clean space"
+        ]
         
-        for i in range(min(num_variations, 8)):
-            # Устанавливаем seed для воспроизводимости
-            generator = torch.Generator(device=self.device).manual_seed(seeds[i])
-            
+        for i in range(min(num_variations, len(prompts))):
             variation = self.pipe(
-                prompt="modern interior design, professional photography",
+                prompt=prompts[i],
                 image=image,
-                strength=strengths[i],
-                num_inference_steps=25,
-                guidance_scale=7.5,
-                generator=generator
+                strength=0.5 + (i * 0.1),
+                num_inference_steps=30,
+                guidance_scale=7.5
             ).images[0]
-            
             variations.append(variation)
-        
+            
         return variations
     
     def create_hdr_lighting(self, image: Image.Image, intensity: float = 0.3) -> Image.Image:
-        """Создание HDR освещения"""
+        """Улучшение освещения (HDR эффект)"""
         # Конвертируем в numpy
-        img_np = np.array(image).astype(np.float32) / 255.0
+        img_array = np.array(image).astype(np.float32)
         
         # Создаем HDR эффект
         # Увеличиваем контраст в светлых областях
-        highlights = np.where(img_np > 0.7, img_np, 0)
-        highlights = highlights ** 0.8
+        bright_mask = img_array > 200
+        img_array[bright_mask] = img_array[bright_mask] * (1 + intensity)
         
-        # Углубляем тени
-        shadows = np.where(img_np < 0.3, img_np, 0)
-        shadows = shadows ** 1.2
+        # Усиливаем тени
+        dark_mask = img_array < 50
+        img_array[dark_mask] = img_array[dark_mask] * (1 - intensity * 0.5)
         
-        # Комбинируем
-        hdr = img_np + (highlights - img_np) * intensity + (shadows - img_np) * intensity * 0.5
+        # Нормализация
+        img_array = np.clip(img_array, 0, 255).astype(np.uint8)
         
-        # Нормализуем
-        hdr = np.clip(hdr, 0, 1)
+        # Применяем дополнительные фильтры
+        enhanced = Image.fromarray(img_array)
+        enhanced = enhanced.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
         
-        # Добавляем легкое свечение
-        blur = cv2.GaussianBlur(hdr, (21, 21), 0)
-        glow = np.where(blur > 0.8, blur, 0)
-        hdr = hdr + glow * intensity * 0.3
+        # Настройка яркости и контраста
+        enhancer = ImageEnhance.Contrast(enhanced)
+        enhanced = enhancer.enhance(1.1)
         
-        # Финальная нормализация
-        hdr = np.clip(hdr, 0, 1)
-        hdr = (hdr * 255).astype(np.uint8)
-        
-        return Image.fromarray(hdr)
+        return enhanced
     
     def enhance_details(self, image: Image.Image) -> Image.Image:
         """Улучшение деталей изображения"""
-        # Конвертируем в numpy
-        img_np = np.array(image)
-        
         # Увеличиваем резкость
-        kernel = np.array([[-1,-1,-1],
-                          [-1, 9,-1],
-                          [-1,-1,-1]])
-        sharpened = cv2.filter2D(img_np, -1, kernel)
-        
-        # Смешиваем с оригиналом
-        result = cv2.addWeighted(img_np, 0.7, sharpened, 0.3, 0)
+        enhanced = image.filter(ImageFilter.SHARPEN)
+        enhanced = enhanced.filter(ImageFilter.UnsharpMask(radius=2, percent=200, threshold=3))
         
         # Улучшаем цвета
-        lab = cv2.cvtColor(result, cv2.COLOR_RGB2LAB)
-        l, a, b = cv2.split(lab)
-        
-        # Улучшаем яркость
-        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
-        l = clahe.apply(l)
+        color_enhancer = ImageEnhance.Color(enhanced)
+        enhanced = color_enhancer.enhance(1.1)
         
-        # Собираем обратно
-        enhanced = cv2.merge([l, a, b])
-        enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB)
+        # Увеличиваем четкость
+        sharpness_enhancer = ImageEnhance.Sharpness(enhanced)
+        enhanced = sharpness_enhancer.enhance(1.2)
         
-        return Image.fromarray(enhanced)
+        return enhanced
     
     def change_element(self, image: Image.Image, element: str, value: str, 
                       strength: float = 0.7) -> Image.Image:
         """Изменение отдельного элемента интерьера"""
         from design_styles import ROOM_ELEMENTS
         
-        element_data = ROOM_ELEMENTS.get(element, {})
-        prompt_add = element_data.get("prompt_add", element.lower())
+        if element not in ROOM_ELEMENTS:
+            raise ValueError(f"Неизвестный элемент: {element}")
+            
+        element_info = ROOM_ELEMENTS[element]
         
-        # Формируем промпт для изменения
-        prompt = f"interior with {value} {prompt_add}, high quality, professional photo"
-        negative = "low quality, distorted, unrealistic"
+        # Создаем промпт для изменения
+        prompt = f"interior design, {element_info['prompt_add']}, {value}"
+        negative_prompt = "blurry, distorted"
         
         # Применяем изменения
         result = self.pipe(
             prompt=prompt,
-            negative_prompt=negative,
+            negative_prompt=negative_prompt,
             image=image,
             strength=strength,
-            num_inference_steps=30,
+            num_inference_steps=40,
             guidance_scale=8.0
         ).images[0]
         
@@ -245,64 +248,77 @@ class InteriorDesignerPro:
     
     def create_style_comparison(self, image: Image.Image, styles: List[str], 
                                room_type: str = "living room") -> Image.Image:
-        """Создание сравнения нескольких стилей"""
+        """Создание сравнения стилей"""
         styled_images = []
         
         for style in styles:
-            styled = self.apply_style_pro(image, style, room_type, strength=0.75, quality="fast")
-            styled_images.append((styled, style))
-        
-        # Создаем сетку для сравнения
-        return self._create_comparison_grid(styled_images, original=image)
-    
-    def _create_comparison_grid(self, styled_images: List[tuple], original: Optional[Image.Image] = None) -> Image.Image:
-        """Создание сетки для сравнения стилей"""
-        images = [img for img, _ in styled_images]
-        titles = [style for _, style in styled_images]
-        
-        # Определяем размер сетки
-        n = len(images)
-        if n <= 2:
-            rows, cols = 1, 2
-        elif n <= 4:
-            rows, cols = 2, 2
-        elif n <= 6:
-            rows, cols = 2, 3
+            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
+            
+        # Используем метод из процессора
+        if hasattr(self, '_create_comparison_grid'):
+            return self._create_comparison_grid(styled_images)
         else:
-            rows, cols = 3, 3
-        
-        # Размер одного изображения
-        img_width = 512
-        img_height = 512
-        padding = 20
-        title_height = 40
-        
-        # Создаем холст
-        grid_width = cols * img_width + (cols + 1) * padding
-        grid_height = rows * (img_height + title_height) + (rows + 1) * padding
+            # Простая сетка
+            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 None
+            
+        images = [img for img, _ in images_with_titles]
+        titles = [title for _, title in images_with_titles]
         
-        if original is not None:
-            # Добавляем место для оригинала сверху
-            grid_height += img_height + title_height + padding
+        # Размеры
+        img_width, img_height = images[0].size
+        cols = min(3, len(images))
+        rows = (len(images) + cols - 1) // cols
         
-        grid = Image.new('RGB', (grid_width, grid_height), color='white')
+        # Создаем сетку
+        grid_width = img_width * cols
+        grid_height = img_height * rows
+        grid = Image.new('RGB', (grid_width, grid_height), 'white')
         
-        # Добавляем изображения на сетку
-        for idx, (img, title) in enumerate(zip(images[:n], titles[:n])):
+        # Размещаем изображения
+        for idx, (img, title) in enumerate(images_with_titles):
             row = idx // cols
             col = idx % cols
+            x = col * img_width
+            y = row * img_height
+            grid.paste(img, (x, y))
             
-            # Изменяем размер
-            img_resized = img.resize((img_width, img_height), Image.Resampling.LANCZOS)
-            
-            # Позиция
-            x = col * (img_width + padding) + padding
-            y = row * (img_height + title_height + padding) + padding
-            
-            if original is not None:
-                y += img_height + title_height + padding
-            
-            # Вставляем изображение
-            grid.paste(img_resized, (x, y))
-        
         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:
+        """Генерация маски на основе текстового описания"""
+        # Простая реализация - создаем маску в центре
+        width, height = image.size
+        mask = Image.new('L', (width, height), 0)
+        draw = ImageDraw.Draw(mask)
+        
+        # Размер области зависит от precision
+        margin = int(min(width, height) * (1 - precision) / 2)
+        
+        # Рисуем белую область в центре
+        draw.ellipse([margin, margin, width-margin, height-margin], fill=255)
+        
+        # Размываем для плавных краев
+        mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
+        
+        return mask