Update models.py

models.py

@@ -3,35 +3,30 @@ from diffusers import (
     StableDiffusionImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
     DDIMScheduler,
-    PNDMScheduler,
-    EulerDiscreteScheduler
+    EulerDiscreteScheduler,
+    PNDMScheduler
 )
 from PIL import Image, ImageFilter, ImageEnhance
 import numpy as np
-from typing import List, Union, Optional
-import cv2
+from typing import List, Optional, Tuple
 
 class InteriorDesignerPro:
     def __init__(self):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        """Инициализация AI дизайнера"""
+        self.device = torch.device("cuda")  # ТОЛЬКО GPU!
         self.model_name = "stabilityai/stable-diffusion-2-1"
         
         # Определяем мощность GPU
-        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 'Нет'}")
-            
-        # Основная модель
+        gpu_name = torch.cuda.get_device_name(0)
+        self.is_powerful_gpu = any(x in gpu_name.lower() for x in ['h100', 'a100', 'h200', 'a6000'])
+        
+        print(f"🚀 Initializing on GPU: {gpu_name}")
+        print(f"💪 Powerful GPU detected: {self.is_powerful_gpu}")
+        
+        # Основная модель для дизайна
         self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
             self.model_name,
-            torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
+            torch_dtype=torch.float16,
             safety_checker=None,
             requires_safety_checker=False
         ).to(self.device)
@@ -40,68 +35,58 @@ class InteriorDesignerPro:
         try:
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting",
-                torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
+                torch_dtype=torch.float16,
                 safety_checker=None,
                 requires_safety_checker=False,
                 local_files_only=False,
                 resume_download=True
             ).to(self.device)
-            print("✅ Inpainting модель загружена")
+            print("✅ Inpainting model loaded")
         except Exception as e:
-            print(f"⚠️ Не удалось загрузить inpainting модель: {e}")
-            print("Используем основную модель для inpainting")
+            print(f"⚠️ Failed to load inpainting model: {e}")
+            print("Using img2img as fallback for object removal")
             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:
-                    print("⚠️ xFormers недоступен")
-                    
+        self.pipe.enable_attention_slicing()
+        if hasattr(self.pipe, 'enable_xformers_memory_efficient_attention'):
+            try:
+                self.pipe.enable_xformers_memory_efficient_attention()
+                print("✅ xFormers optimization enabled")
+            except:
+                pass
+        
         # Настройка планировщиков
         self.schedulers = {
-            "ddim": DDIMScheduler.from_config(self.pipe.scheduler.config),
-            "pndm": PNDMScheduler.from_config(self.pipe.scheduler.config),
-            "euler": EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
+            "fast": EulerDiscreteScheduler.from_config(self.pipe.scheduler.config),
+            "balanced": DDIMScheduler.from_config(self.pipe.scheduler.config),
+            "quality": PNDMScheduler.from_config(self.pipe.scheduler.config)
         }
-        
-        # Удаление объектов
-        self.object_remover = ObjectRemover(self.device)
-        
+    
     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
+        """Применение стиля к интерьеру"""
+        from design_styles import DESIGN_STYLES, get_detailed_prompt
         
-        if style not in DESIGN_STYLES:
-            style = "Современный минимализм"
-            
-        style_info = DESIGN_STYLES[style]
-        base_prompt = style_info.get("prompt", "")
+        # Подготовка изображения
+        image = self._prepare_image(image)
+        
+        # Получаем промпт для стиля
+        style_info = DESIGN_STYLES.get(style, DESIGN_STYLES["Современный минимализм"])
+        prompt = get_detailed_prompt(style, room_type)
         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"
-            
-        # Настройки качества
-        quality_settings = {
-            "fast": {"steps": 20, "guidance": 7.5, "scheduler": "euler"},
-            "balanced": {"steps": 35, "guidance": 8.5, "scheduler": "ddim"},
-            "ultra": {"steps": 50, "guidance": 10, "scheduler": "ddim"}
+        # Выбираем планировщик
+        self.pipe.scheduler = self.schedulers.get(quality, self.schedulers["balanced"])
+        
+        # Параметры генерации
+        params = {
+            "fast": {"steps": 20, "guidance": 7.5},
+            "balanced": {"steps": 35, "guidance": 8.5},
+            "ultra": {"steps": 50, "guidance": 10}
         }
         
-        settings = quality_settings.get(quality, quality_settings["balanced"])
-        self.pipe.scheduler = self.schedulers[settings["scheduler"]]
+        settings = params.get(quality, params["balanced"])
         
         # Генерация
         result = self.pipe(
@@ -114,90 +99,133 @@ class InteriorDesignerPro:
         ).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 метод через img2img
+            return self._remove_objects_fallback(image, mask)
+        
+        # Подготовка
+        image = self._prepare_image(image, size=(512, 512))
+        mask = self._prepare_mask(mask, size=(512, 512))
+        
+        # Промпт для заполнения
+        if not inpaint_prompt:
+            inpaint_prompt = "empty room interior, clean walls, seamless texture, high quality"
+        
+        # Inpainting
+        result = self.inpaint_pipe(
+            prompt=inpaint_prompt,
+            negative_prompt="furniture, objects, people, clutter",
+            image=image,
+            mask_image=mask,
+            num_inference_steps=50,
+            guidance_scale=7.5,
+            strength=0.99
+        ).images[0]
+        
+        return result
+    
+    def _remove_objects_fallback(self, image: Image.Image, mask: Image.Image) -> Image.Image:
+        """Альтернативный метод удаления объектов"""
+        # Размываем область под маской
+        import cv2
+        
+        img_np = np.array(image)
+        mask_np = np.array(mask.convert('L'))
+        
+        # Расширяем маску
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
+        mask_np = cv2.dilate(mask_np, kernel, iterations=2)
         
+        # Inpaint через OpenCV
+        result_np = cv2.inpaint(img_np, mask_np, 7, cv2.INPAINT_TELEA)
+        result = Image.fromarray(result_np)
+        
+        # Улучшаем через img2img
+        enhanced = self.pipe(
+            prompt="clean empty interior space, seamless walls",
+            negative_prompt="objects, furniture",
+            image=result,
+            strength=0.3,
+            num_inference_steps=20,
+            guidance_scale=5
+        ).images[0]
+        
+        return enhanced
+    
     def create_variations(self, image: Image.Image, num_variations: int = 4,
-                         strength: float = 0.5) -> List[Image.Image]:
-        """Создание вариаций изображения"""
+                         variation_strength: float = 0.3) -> List[Image.Image]:
+        """Создание вариаций дизайна"""
         variations = []
+        base_prompt = "interior design, professional photo, high quality"
         
-        prompts = [
-            "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))):
-            var = self.pipe(
-                prompt=prompts[i],
+        for i in range(num_variations):
+            # Немного меняем seed для разнообразия
+            generator = torch.Generator(device=self.device).manual_seed(42 + i)
+            
+            result = self.pipe(
+                prompt=base_prompt,
                 image=image,
-                strength=strength,
-                num_inference_steps=25,
-                guidance_scale=7.5
+                strength=variation_strength,
+                num_inference_steps=30,
+                guidance_scale=7.5,
+                generator=generator
             ).images[0]
-            variations.append(var)
             
-        return variations
+            variations.append(result)
         
+        return variations
+    
     def create_hdr_lighting(self, image: Image.Image, intensity: float = 0.3) -> Image.Image:
-        """Улучшение освещения (HDR эффект)"""
+        """Создание HDR освещения"""
         # Конвертируем в numpy
-        img_array = np.array(image)
+        img_array = np.array(image).astype(np.float32) / 255.0
         
-        # Создаем HDR эффект
-        # Светлая версия
-        light = cv2.convertScaleAbs(img_array, alpha=1.2, beta=30)
-        # Темная версия  
-        dark = cv2.convertScaleAbs(img_array, alpha=0.8, beta=-30)
-        
-        # Комбинируем
-        hdr = cv2.addWeighted(img_array, 1-intensity, light, intensity/2, 0)
-        hdr = cv2.addWeighted(hdr, 1, dark, intensity/2, 0)
+        # Создаем несколько экспозиций
+        exposures = [
+            img_array * 0.5,  # Недоэкспонированная
+            img_array,        # Нормальная
+            np.clip(img_array * 1.5, 0, 1)  # Переэкспонированная
+        ]
         
-        # Улучшаем контраст
-        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)
+        # Простой tone mapping
+        hdr = np.mean(exposures, axis=0)
         
-        return Image.fromarray(result)
+        # Увеличиваем контраст
+        hdr = (hdr - 0.5) * (1 + intensity) + 0.5
+        hdr = np.clip(hdr, 0, 1)
         
+        # Обратно в PIL
+        return Image.fromarray((hdr * 255).astype(np.uint8))
+    
     def enhance_details(self, image: Image.Image) -> Image.Image:
         """Улучшение деталей изображения"""
         # Увеличиваем резкость
         enhancer = ImageEnhance.Sharpness(image)
-        sharp = enhancer.enhance(1.5)
+        image = enhancer.enhance(1.5)
         
-        # Улучшаем детали
-        detail_filter = ImageFilter.DETAIL
-        detailed = sharp.filter(detail_filter)
+        # Немного увеличиваем контраст
+        enhancer = ImageEnhance.Contrast(image)
+        image = enhancer.enhance(1.1)
         
-        # Легкий 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)
+        # Применяем unsharp mask
+        image = image.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
         
+        return image
+    
     def change_element(self, image: Image.Image, element: str, value: str, 
                       strength: float = 0.5) -> Image.Image:
-        """Изменение конкретного элемента интерьера"""
+        """Изменение отдельного элемента интерьера"""
         from design_styles import ROOM_ELEMENTS
         
-        if element not in ROOM_ELEMENTS:
-            return image
-            
-        element_info = ROOM_ELEMENTS[element]
-        prompt_add = element_info.get("prompt_add", "")
+        element_info = ROOM_ELEMENTS.get(element, {})
+        prompt_add = element_info.get("prompt_add", element.lower())
         
-        prompt = f"{element} {value}, {prompt_add}, interior design"
-        negative = "bad quality, distorted"
+        prompt = f"interior with {value} {prompt_add}, professional photo"
+        negative = f"old {element}, damaged, ugly"
         
         result = self.pipe(
             prompt=prompt,
@@ -205,124 +233,64 @@ class InteriorDesignerPro:
             image=image,
             strength=strength,
             num_inference_steps=30,
-            guidance_scale=8.0
+            guidance_scale=8
         ).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:
+    
+    def create_style_comparison(self, image: Image.Image, styles: List[str]) -> Image.Image:
         """Создание сравнения нескольких стилей"""
         styled_images = []
         
-        # Генерируем для каждого стиля
         for style in styles:
-            styled = self.apply_style_pro(
-                image, style, room_type, 
-                strength=0.75, 
-                quality="fast"  # Быстрый режим для множественной генерации
-            )
+            styled = self.apply_style_pro(image, style, "living room", strength=0.75, quality="fast")
             styled_images.append(styled)
-            
-        # Создаем сетку
+        
+        # Создаем сетку через метод который добавим динамически
         return self._create_comparison_grid(styled_images, styles)
+    
+    def _prepare_image(self, image: Image.Image, size: tuple = (768, 768)) -> Image.Image:
+        """Подготовка изображения"""
+        # Конвертируем в RGB
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Изменяем размер сохраняя пропорции
+        image.thumbnail(size, Image.Resampling.LANCZOS)
+        
+        # Дополняем до нужного размера
+        new_image = Image.new('RGB', size, (255, 255, 255))
+        paste_x = (size[0] - image.width) // 2
+        paste_y = (size[1] - image.height) // 2
+        new_image.paste(image, (paste_x, paste_y))
+        
+        return new_image
+    
+    def _prepare_mask(self, mask: Image.Image, size: tuple = (512, 512)) -> Image.Image:
+        """Подготовка маски"""
+        # Конвертируем в L (grayscale)
+        if mask.mode != 'L':
+            mask = mask.convert('L')
+        
+        # Изменяем размер
+        mask = mask.resize(size, Image.Resampling.LANCZOS)
+        
+        # Усиливаем контраст (четкие границы)
+        mask = mask.point(lambda x: 255 if x > 128 else 0)
         
-    def _create_comparison_grid(self, images: List[Image.Image], 
-                               titles: List[str]) -> Image.Image:
-        """Создание сетки из изображений"""
-        if not images:
-            return None
-            
-        # Определяем размер сетки
-        n = len(images)
-        cols = min(3, n)
-        rows = (n + cols - 1) // cols
-        
-        # Размер одного изображения
-        img_width, img_height = images[0].size
-        
-        # Создаем холст
-        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(zip(images, titles)):
-            row = idx // cols
-            col = idx % cols
-            x = col * img_width
-            y = row * img_height
-            grid.paste(img, (x, y))
-            
-        return grid
-
+        return mask
 
+# Класс для удаления объектов
 class ObjectRemover:
-    """Класс для удаления объектов"""
-    def __init__(self, device):
-        self.device = device
-        
-    def generate_mask_from_text(self, image: Image.Image, text: str, 
+    def __init__(self):
+        self.device = torch.device("cuda")  # ТОЛЬКО GPU!
+    
+    def generate_mask_from_text(self, image: Image.Image, text_description: str, 
                                precision: float = 0.3) -> Image.Image:
         """Генерация маски на основе текстового описания"""
         # Простая реализация - создаем маску в центре
+        # В реальности здесь должен быть CLIP/SAM для поиска объектов
+        
         width, height = image.size
         mask = Image.new('L', (width, height), 0)
         
@@ -334,12 +302,13 @@ class ObjectRemover:
         margin_x = int(width * (0.5 - precision))
         margin_y = int(height * (0.5 - precision))
         
-        draw.ellipse(
+        # Рисуем белую область
+        draw.rectangle(
             [margin_x, margin_y, width - margin_x, height - margin_y],
             fill=255
         )
         
-        # Размываем маску
+        # Размываем края для плавного перехода
         mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
         
         return mask