Update models.py

models.py

@@ -3,312 +3,226 @@ from diffusers import (
     StableDiffusionImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
     DDIMScheduler,
-    EulerDiscreteScheduler,
-    PNDMScheduler
+    PNDMScheduler,
+    EulerDiscreteScheduler
 )
 from PIL import Image, ImageFilter, ImageEnhance
 import numpy as np
-from typing import List, Optional, Tuple
+import cv2
 
 class InteriorDesignerPro:
     def __init__(self):
-        """Инициализация AI дизайнера"""
-        self.device = torch.device("cuda")  # ТОЛЬКО GPU!
-        self.model_name = "stabilityai/stable-diffusion-2-1"
-        
-        # Определяем мощность GPU
+        self.device = torch.device("cuda")
+        self.model_name = "RealVisXL V4.0"
+
+        # Проверка GPU
         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.is_powerful_gpu = any(gpu in gpu_name for gpu in ['A100', 'H100', 'RTX 4090', 'RTX 3090'])
+
+        # Основная модель - RealVis V4 для фотореалистичных интерьеров
+        print(f"Loading {self.model_name} on {gpu_name}...")
         self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-            self.model_name,
+            "SG161222/RealVisXL_V4.0",
             torch_dtype=torch.float16,
             safety_checker=None,
-            requires_safety_checker=False
+            requires_safety_checker=False,
+            local_files_only=False
         ).to(self.device)
-        
-        # Модель для inpainting (удаление объектов)
+
+        # Настройка scheduler для лучшего качества
+        self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
+
+        # Inpainting модель для удаления объектов
         try:
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
-                "runwayml/stable-diffusion-inpainting",
+                "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")
+            print("Inpainting model loaded successfully")
         except Exception as e:
-            print(f"⚠️ Failed to load inpainting model: {e}")
+            print(f"Warning: Could not load inpainting model: {e}")
             print("Using img2img as fallback for object removal")
             self.inpaint_pipe = None
-        
-        # Оптимизация памяти
-        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 = {
-            "fast": EulerDiscreteScheduler.from_config(self.pipe.scheduler.config),
-            "balanced": DDIMScheduler.from_config(self.pipe.scheduler.config),
-            "quality": PNDMScheduler.from_config(self.pipe.scheduler.config)
-        }
-    
-    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, get_detailed_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", "")
-        
-        # Выбираем планировщик
-        self.pipe.scheduler = self.schedulers.get(quality, self.schedulers["balanced"])
-        
-        # Параметры генерации
-        params = {
+            
+    def apply_style_pro(self, image, style_name, room_type, strength=0.75, quality="balanced"):
+        """Применение стиля к изображению с учетом качества"""
+        from design_styles import DESIGN_STYLES
+
+        style = DESIGN_STYLES.get(style_name, DESIGN_STYLES["Современный минимализм"])
+
+        # Настройки качества
+        quality_settings = {
             "fast": {"steps": 20, "guidance": 7.5},
             "balanced": {"steps": 35, "guidance": 8.5},
             "ultra": {"steps": 50, "guidance": 10}
         }
-        
-        settings = params.get(quality, params["balanced"])
-        
+
+        settings = quality_settings.get(quality, quality_settings["balanced"])
+
+        # Генерация промпта с учетом комнаты
+        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"
+
         # Генерация
         result = self.pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
+            prompt=full_prompt,
+            negative_prompt=style.get("negative", "low quality, blurry"),
             image=image,
             strength=strength,
             num_inference_steps=settings["steps"],
             guidance_scale=settings["guidance"]
         ).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,
-                         variation_strength: float = 0.3) -> List[Image.Image]:
+
+    def create_variations(self, image, num_variations=4):
         """Создание вариаций дизайна"""
         variations = []
-        base_prompt = "interior design, professional photo, high quality"
-        
+        base_seed = torch.randint(0, 1000000, (1,)).item()
+
         for i in range(num_variations):
-            # Немного меняем seed для разнообразия
-            generator = torch.Generator(device=self.device).manual_seed(42 + i)
-            
-            result = self.pipe(
-                prompt=base_prompt,
+            torch.manual_seed(base_seed + i)
+
+            var = self.pipe(
+                prompt="interior design variation, same style, different details",
                 image=image,
-                strength=variation_strength,
+                strength=0.4 + (i * 0.05),
                 num_inference_steps=30,
-                guidance_scale=7.5,
-                generator=generator
+                guidance_scale=7.5
             ).images[0]
-            
-            variations.append(result)
-        
+
+            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) / 255.0
-        
-        # Создаем несколько экспозиций
-        exposures = [
-            img_array * 0.5,  # Недоэкспонированная
-            img_array,        # Нормальная
-            np.clip(img_array * 1.5, 0, 1)  # Переэкспонированная
-        ]
-        
-        # Простой tone mapping
-        hdr = np.mean(exposures, axis=0)
         
-        # Увеличиваем контраст
-        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:
+    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)
-        image = enhancer.enhance(1.5)
-        
+        sharp = enhancer.enhance(1.5)
+
         # Немного увеличиваем контраст
-        enhancer = ImageEnhance.Contrast(image)
-        image = enhancer.enhance(1.1)
-        
-        # Применяем 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:
+        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
-        
+
         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,
             negative_prompt=negative,
             image=image,
             strength=strength,
-            num_inference_steps=30,
-            guidance_scale=8
+            num_inference_steps=40,
+            guidance_scale=8.0
         ).images[0]
-        
+
         return result
-    
-    def create_style_comparison(self, image: Image.Image, styles: List[str]) -> Image.Image:
-        """Создание сравнения нескольких стилей"""
-        styled_images = []
         
+    def create_style_comparison(self, image, styles, quality="fast"):
+        """Создание сравнения стилей"""
+        results = []
+
+        # Настройки для быстрой генерации
+        steps = 20 if quality == "fast" else 35
+
         for style in styles:
-            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)
-        
-        return mask
+            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):
-        self.device = torch.device("cuda")  # ТОЛЬКО GPU!
-    
-    def generate_mask_from_text(self, image: Image.Image, text_description: str, 
-                               precision: float = 0.3) -> Image.Image:
+    """Класс для удаления объектов"""
+
+    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)
+
+        # Используем 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.99,
+            num_inference_steps=50,
+            guidance_scale=7.5
+        ).images[0]
+
+        return result
+
+    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 для поиска объектов
-        
+        # Простая маска в центре (заглушка)
+        # В реальности тут должен быть CLIP или SAM
         width, height = image.size
         mask = Image.new('L', (width, height), 0)
-        
-        # Создаем примерную область в центре
-        from PIL import ImageDraw
+
+        # Создаем маску в центре
+        center_x, center_y = width // 2, height // 2
+        radius = int(min(width, height) * precision)
+
+        # Рисуем круг
+        import ImageDraw
         draw = ImageDraw.Draw(mask)
-        
-        # Размер области зависит от precision
-        margin_x = int(width * (0.5 - precision))
-        margin_y = int(height * (0.5 - precision))
-        
-        # Рисуем белую область
-        draw.rectangle(
-            [margin_x, margin_y, width - margin_x, height - margin_y],
-            fill=255
-        )
-        
-        # Размываем края для плавного перехода
-        mask = mask.filter(ImageFilter.GaussianBlur(radius=20))
-        
+        draw.ellipse([center_x - radius, center_y - radius, 
+                      center_x + radius, center_y + radius], fill=255)
+
         return mask