Update models.py

models.py

@@ -1,384 +1,252 @@
 import torch
 from diffusers import (
-    StableDiffusionXLImg2ImgPipeline,
     StableDiffusionImg2ImgPipeline,
-    AutoPipelineForImage2Image,
-    DPMSolverMultistepScheduler,
-    EulerAncestralDiscreteScheduler,
-    AutoencoderKL
+    StableDiffusionInpaintPipeline,
+    DDIMScheduler,
+    PNDMScheduler,
+    EulerDiscreteScheduler
 )
-from PIL import Image
+from PIL import Image, ImageEnhance, ImageFilter
 import numpy as np
+from typing import Optional, List, Union
 
 class InteriorDesignerPro:
+    """AI модель для профессионального дизайна интерьеров"""
+    
     def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
-
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.dtype = torch.float16 if self.device.type == 'cuda' else torch.float32
+        
         # Определяем мощность GPU
         if torch.cuda.is_available():
-            gpu_name = torch.cuda.get_device_name()
-            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
-            self.is_powerful_gpu = gpu_memory > 20  # Больше 20GB VRAM
+            gpu_name = torch.cuda.get_device_name(0).lower()
+            self.is_powerful_gpu = any(x in gpu_name for x in ['a100', 'h100', 'rtx 4090', 'rtx 3090', 'v100', 'h200'])
         else:
             self.is_powerful_gpu = False
-            gpu_name = "CPU"
-            gpu_memory = 0
-
-        print(f"GPU: {gpu_name}")
-        print(f"VRAM: {gpu_memory:.1f}GB")
-        print(f"Powerful GPU mode: {self.is_powerful_gpu}")
-
-        self.init_models()
-
-    def init_models(self):
-        """Инициализация моделей с safetensors"""
-        self.pipe = None
-        self.turbo_pipe = None
-
-        # Топовые модели для интерьеров (проверено на safetensors!)
-        models_to_try = [
-            {
-                "name": "RealVisXL V4.0 (ЛУЧШАЯ для интерьеров)",
-                "model_id": "SG161222/RealVisXL_V4.0",
-                "type": "sdxl"
-            },
-            {
-                "name": "Juggernaut XL v9 (Фотореализм)",
-                "model_id": "RunDiffusion/Juggernaut-XL-v9",
-                "type": "sdxl"
-            },
-            {
-                "name": "DreamShaper XL Turbo (Быстрый)",
-                "model_id": "Lykon/dreamshaper-xl-turbo",
-                "type": "sdxl"
-            },
-            {
-                "name": "SDXL Base + Refiner",
-                "model_id": "stabilityai/stable-diffusion-xl-base-1.0",
-                "type": "sdxl"
-            },
-            {
-                "name": "Realistic Vision V6.0 B1",
-                "model_id": "SG161222/Realistic_Vision_V6.0_B1_noVAE",
-                "type": "sd15"
-            }
-        ]
-
-        # VAE для улучшения цветов (опционально)
-        vae = None
-        try:
-            vae = AutoencoderKL.from_pretrained(
-                "madebyollin/sdxl-vae-fp16-fix",
-                torch_dtype=self.dtype,
-                use_safetensors=True
-            ).to(self.device)
-            print("✅ VAE для улучшения цветов загружен")
-        except:
-            print("⚠️ VAE не загружен, используем встроенный")
-
-        # Пытаемся загрузить модели по очереди
-        for model_info in models_to_try:
+            
+        # Основная модель для дизайна
+        model_id = "runwayml/stable-diffusion-v1-5"
+        self.model_name = "SD 1.5 Professional"
+        
+        print(f"🚀 Loading {self.model_name} on {self.device}")
+        print(f"💪 Powerful GPU detected: {self.is_powerful_gpu}")
+        
+        # Загружаем пайплайны
+        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            model_id,
+            torch_dtype=self.dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(self.device)
+        
+        # Пайплайн для inpainting
+        self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            torch_dtype=self.dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(self.device)
+        
+        # Оптимизация памяти
+        if self.device.type == 'cuda':
+            self.pipe.enable_attention_slicing()
+            self.inpaint_pipe.enable_attention_slicing()
+            
+            # xFormers если доступно
             try:
-                print(f"\nПытаюсь загрузить: {model_info['name']}")
-
-                if model_info['type'] == 'sdxl':
-                    # Загрузка SDXL модели
-                    self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
-                        model_info['model_id'],
-                        vae=vae,  # Используем улучшенный VAE если есть
-                        torch_dtype=self.dtype,
-                        use_safetensors=True,
-                        variant="fp16" if self.device == "cuda" else None,
-                        safety_checker=None
-                    ).to(self.device)
-
-                else:
-                    # Загрузка SD 1.5 модели
-                    self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-                        model_info['model_id'],
-                        torch_dtype=self.dtype,
-                        use_safetensors=True,
-                        safety_checker=None,
-                        requires_safety_checker=False
-                    ).to(self.device)
-
-                print(f"✅ Успешно загружена модель: {model_info['name']}")
-                self.model_type = model_info['type']
-                self.model_name = model_info['name']
-                break
-
-            except Exception as e:
-                print(f"❌ Не удалось загрузить {model_info['name']}: {str(e)[:100]}...")
-                continue
-
-        if self.pipe is None:
-            # Последняя попытка - минимальная модель
-            try:
-                print("\nПоследняя попытка - загружаю минимальную SD 1.5...")
-                self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-                    "runwayml/stable-diffusion-v1-5",
-                    torch_dtype=self.dtype,
-                    safety_checker=None,
-                    requires_safety_checker=False,
-                    use_safetensors=True
-                ).to(self.device)
-                self.model_type = "sd15"
-                self.model_name = "SD 1.5 Base"
-                print("✅ Загружена базовая SD 1.5")
-            except:
-                raise Exception("Не удалось загрузить ни одну модель!")
-
-        # Настройка scheduler для лучшего качества
-        try:
-            if self.model_type == "sdxl":
-                self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
-                    self.pipe.scheduler.config
-                )
-                print("✅ Настроен EulerAncestral scheduler")
-        except:
-            print("⚠️ Используется стандартный scheduler")
-
-        # Оптимизации
-        try:
-            if self.device == "cuda":
                 self.pipe.enable_xformers_memory_efficient_attention()
-                self.pipe.enable_vae_slicing()
-                if hasattr(self.pipe, 'enable_vae_tiling'):
-                    self.pipe.enable_vae_tiling()
-                print("✅ Включены оптимизации памяти")
-        except:
-            print("⚠️ Некоторые оптимизации недоступны")
-
-        # Загружаем turbo для preview (если есть память)
-        if self.is_powerful_gpu:
-            try:
-                print("\nЗагружаю SDXL Turbo для быстрого preview...")
-                self.turbo_pipe = AutoPipelineForImage2Image.from_pretrained(
-                    "stabilityai/sdxl-turbo",
-                    torch_dtype=self.dtype,
-                    variant="fp16",
-                    use_safetensors=True
-                ).to(self.device)
-                print("✅ SDXL Turbo загружен")
+                self.inpaint_pipe.enable_xformers_memory_efficient_attention()
+                print("✅ xFormers optimization enabled")
             except:
-                print("⚠️ Не удалось загрузить Turbo модель")
-                self.turbo_pipe = None
-
-    def apply_style_pro(self, image, style_name, room_type="living room", 
-                       strength=0.75, quality="balanced"):
-        """Применение стиля дизайна с максимальным качеством"""
-        from design_styles import DESIGN_STYLES
-
-        if style_name not in DESIGN_STYLES:
-            return image
-
-        style = DESIGN_STYLES[style_name]
-
-        # Настройки качества
-        quality_settings = {
-            "fast": {
-                "steps": 4 if self.turbo_pipe else 20,
-                "cfg": 2.0 if self.turbo_pipe else 7.5,
-                "resolution": 768 if self.model_type == "sd15" else 1024
-            },
-            "balanced": {
-                "steps": 30,
-                "cfg": 7.5,
-                "resolution": 1024 if self.model_type == "sd15" else 1536
-            },
-            "ultra": {
-                "steps": 50,
-                "cfg": 8.0,
-                "resolution": 1024 if self.model_type == "sd15" else 2048
-            }
-        }
-
-        settings = quality_settings.get(quality, quality_settings["balanced"])
-
-        # Подготовка изображения
-        target_size = settings["resolution"]
-        if not self.is_powerful_gpu:
-            target_size = min(target_size, 1024)
-
-        # Изменение размера с сохранением пропорций
-        aspect_ratio = image.width / image.height
-        if aspect_ratio > 1:
-            new_width = target_size
-            new_height = int(target_size / aspect_ratio)
-        else:
-            new_height = target_size
-            new_width = int(target_size * aspect_ratio)
-
-        # Размеры кратные 8
-        new_width = (new_width // 8) * 8
-        new_height = (new_height // 8) * 8
-
-        image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
-
-        # Улучшенные промпты для интерьеров
-        if "RealVis" in self.model_name or "Juggernaut" in self.model_name:
-            # Промпты для фотореалистичных моделей
-            quality_tags = {
-                "fast": "RAW photo, high quality",
-                "balanced": "RAW photo, masterpiece, best quality, ultra realistic, 8k",
-                "ultra": "RAW photo, masterpiece, best quality, ultra realistic, 8k uhd, film grain, Fujifilm XT3"
-            }
-            photo_tags = "(interior design photography:1.2), architectural digest, elle decor, "
+                pass
+        
+        # Настройка scheduler для старой версии
+        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        self.inpaint_pipe.scheduler = DDIMScheduler.from_config(self.inpaint_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
+        
+        # Получаем детальный промпт
+        style_info = DESIGN_STYLES.get(style, DESIGN_STYLES["Современный минимализм"])
+        
+        # Базовый промпт + специфика комнаты
+        base_prompt = style_info["prompt"]
+        room_specific = style_info.get("room_specific", {}).get(room_type, "")
+        
+        # Комбинируем промпты
+        if room_specific:
+            full_prompt = f"{base_prompt}, {room_specific}"
         else:
-            # Стандартные промпты
-            quality_tags = {
-                "fast": "high quality",
-                "balanced": "masterpiece, best quality, high resolution",
-                "ultra": "masterpiece, best quality, ultra detailed, 8k uhd"
-            }
-            photo_tags = "interior design photography, professional, "
-
-        lighting_tags = "perfect lighting, natural light, soft shadows, "
-        camera_tags = "35mm photograph, professional color grading, "
-
-        prompt = f"{quality_tags[quality]}, {photo_tags}{lighting_tags}{camera_tags}{room_type} interior, {style['prompt']}"
-
-        # Негативные промпты
-        if "RealVis" in self.model_name:
-            negative_prompt = (
-                f"(worst quality:2), (low quality:2), (normal quality:2), lowres, "
-                f"((monochrome)), ((grayscale)), skin spots, acnes, skin blemishes, "
-                f"bad anatomy, {style['negative']}, watermark, signature"
-            )
-        else:
-            negative_prompt = (
-                f"{style['negative']}, low quality, worst quality, blurry, "
-                f"distorted, deformed, ugly, bad anatomy, watermark, signature, text"
-            )
-
-        # Выбор пайплайна
-        pipe = self.turbo_pipe if (quality == "fast" and self.turbo_pipe) else self.pipe
-
+            full_prompt = f"{base_prompt}, {room_type} interior"
+            
+        # Добавляем усилители качества
+        quality_boost = {
+            "fast": "",
+            "balanced": "high quality, professional photo, ",
+            "ultra": "masterpiece, best quality, ultra detailed, 8k uhd, professional photography, "
+        }
+        
+        full_prompt = quality_boost.get(quality, "") + full_prompt
+        
+        # Параметры генерации для старой версии diffusers
+        params = {
+            "fast": {"num_inference_steps": 20, "guidance_scale": 7.5},
+            "balanced": {"num_inference_steps": 35, "guidance_scale": 8.5},
+            "ultra": {"num_inference_steps": 50, "guidance_scale": 10}
+        }
+        
         # Генерация
-        try:
-            # Параметры для SDXL
-            if self.model_type == "sdxl" and hasattr(pipe, 'text_encoder_2'):
-                result = pipe(
-                    prompt=prompt,
-                    prompt_2=f"luxury {room_type} interior design, {style['prompt']}",
-                    negative_prompt=negative_prompt,
-                    negative_prompt_2=negative_prompt,
-                    image=image,
-                    strength=strength,
-                    num_inference_steps=settings["steps"],
-                    guidance_scale=settings["cfg"],
-                    width=new_width,
-                    height=new_height
-                ).images[0]
-            else:
-                # Параметры для SD 1.5
-                result = pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    image=image,
-                    strength=strength,
-                    num_inference_steps=settings["steps"],
-                    guidance_scale=settings["cfg"]
-                ).images[0]
-
-            return result
-
-        except Exception as e:
-            print(f"Ошибка генерации: {e}")
-            return image
-
-    def create_hdr_lighting(self, image, intensity=0.3):
-        """HDR эффект для освещения"""
-        if "RealVis" in self.model_name:
-            hdr_prompt = (
-                "RAW photo, HDR photography, professional lighting setup, "
-                "golden hour, soft natural light, ray tracing, high dynamic range"
-            )
-        else:
-            hdr_prompt = (
-                "HDR photography, perfect lighting, golden hour, "
-                "soft shadows, ambient occlusion, professional lighting"
-            )
-
-        try:
-            return self.pipe(
-                prompt=hdr_prompt,
-                negative_prompt="flat lighting, harsh shadows, overexposed, underexposed",
-                image=image,
-                strength=intensity,
-                num_inference_steps=20,
-                guidance_scale=7.0
-            ).images[0]
-        except Exception as e:
-            print(f"Ошибка HDR: {e}")
-            return image
-
-    def enhance_details(self, image):
-        """Улучшение деталей и резкости"""
-        if not self.is_powerful_gpu:
-            return image
-
-        try:
-            if "RealVis" in self.model_name:
-                detail_prompt = "RAW photo, ultra sharp, highly detailed, 4k textures, fine details"
-            else:
-                detail_prompt = "ultra sharp, highly detailed, 4k, crisp details, clear focus"
-
-            return self.pipe(
-                prompt=detail_prompt,
-                negative_prompt="blurry, soft focus, out of focus",
-                image=image,
-                strength=0.2,
-                num_inference_steps=15,
-                guidance_scale=7.0
-            ).images[0]
-        except:
-            return image
-
-    def create_variations(self, image, num_variations=4):
-        """Создание вариаций дизайна"""
-        from design_styles import DESIGN_STYLES
+        result = self.pipe(
+            prompt=full_prompt,
+            image=image,
+            strength=strength,
+            negative_prompt=style_info.get("negative", ""),
+            **params[quality]
+        ).images[0]
+        
+        return result
+        
+    def create_variations(self, image: Image.Image, num_variations: int = 4) -> List[Image.Image]:
+        """Создает вариации дизайна"""
         variations = []
-        styles = list(DESIGN_STYLES.keys())[:num_variations]
-
-        # Используем fast режим для вариаций
-        for style in styles:
-            try:
-                var = self.apply_style_pro(image, style, strength=0.7, quality="fast")
-                variations.append(var)
-            except Exception as e:
-                print(f"Ошибка вариации {style}: {e}")
-                variations.append(image)
-
+        base_seed = torch.randint(0, 1000000, (1,)).item()
+        
+        for i in range(num_variations):
+            # Устанавливаем seed для вариации
+            torch.manual_seed(base_seed + i)
+            
+            # Случайный стиль и сила
+            from design_styles import DESIGN_STYLES
+            styles = list(DESIGN_STYLES.keys())
+            random_style = styles[i % len(styles)]
+            
+            strength = 0.6 + (i * 0.05)  # От 0.6 до 0.75
+            
+            variation = self.apply_style_pro(
+                image, 
+                random_style, 
+                "living room",
+                strength=strength,
+                quality="fast"
+            )
+            variations.append(variation)
+            
         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
+        
+        # Применяем tone mapping
+        gamma = 1.0 + intensity
+        img_hdr = np.power(img_array, 1.0 / gamma)
+        
+        # Увеличиваем контраст в светлых областях
+        img_hdr = img_hdr * (1.0 + intensity * 0.5)
+        img_hdr = np.clip(img_hdr, 0, 1)
+        
+        # Обратно в PIL
+        img_hdr = (img_hdr * 255).astype(np.uint8)
+        result = Image.fromarray(img_hdr)
+        
+        # Дополнительная обработка
+        enhancer = ImageEnhance.Color(result)
+        result = enhancer.enhance(1.0 + intensity * 0.3)
+        
+        return result
+        
+    def enhance_details(self, image: Image.Image) -> Image.Image:
+        """Улучшает детализацию изображения"""
+        # Увеличиваем резкость
+        enhanced = image.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))
+        
+        # Улучшаем контраст
+        enhancer = ImageEnhance.Contrast(enhanced)
+        enhanced = enhancer.enhance(1.15)
+        
+        # Насыщенность
+        enhancer = ImageEnhance.Color(enhanced)
+        enhanced = enhancer.enhance(1.1)
+        
+        return enhanced
+        
+    def change_element(self, image: Image.Image, element: str, value: str, 
+                      strength: float = 0.5) -> Image.Image:
+        """Изменяет отдельный элемент интерьера"""
+        from design_styles import ROOM_ELEMENTS
+        
+        element_info = ROOM_ELEMENTS.get(element, {})
+        prompt_add = element_info.get("prompt_add", element)
+        
+        # Формируем промпт
+        prompt = f"interior with {value} {prompt_add}, high quality, detailed"
+        
+        # Применяем изменение
+        result = self.pipe(
+            prompt=prompt,
+            image=image,
+            strength=strength,
+            negative_prompt="low quality, blurry",
+            num_inference_steps=30,
+            guidance_scale=8.0
+        ).images[0]
+        
+        return result
+        
+    def create_style_comparison(self, image: Image.Image, styles: List[str], 
+                               room_type: str = "living room") -> Image.Image:
+        """Создает сравнение разных стилей"""
+        from utils import ImageProcessor
+        processor = ImageProcessor()
+        
+        styled_images = []
+        for style in styles:
+            styled = self.apply_style_pro(
+                image, 
+                style, 
+                room_type, 
+                strength=0.75,
+                quality="fast"
+            )
+            styled_images.append(styled)
+            
+        # Создаем сетку
+        comparison = processor.create_grid(styled_images, titles=styles)
+        return comparison
+
+# Добавляем метод для создания сетки если его нет
+def _create_comparison_grid(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 = 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, titles)):
+        row = idx // cols
+        col = idx % cols
+        x = col * img_width
+        y = row * img_height
+        grid.paste(img, (x, y))
+        
+    return grid
+
+# Патчим метод если его нет
+if not hasattr(InteriorDesignerPro, '_create_comparison_grid'):
+    InteriorDesignerPro._create_comparison_grid = _create_comparison_grid
 
-    def change_element(self, image, element, value, strength=0.5):
-        """Изменение отдельного элемента комнаты"""
-        element_prompts = {
-            "Стены": "wall color and texture",
-            "Пол": "floor material and pattern", 
-            "Освещение": "lighting fixtures and ambiance",
-            "Мебель": "furniture style and arrangement",
-            "Декор": "decorations and accessories"
-        }
-
-        base_element = element_prompts.get(element, element)
-
-        if "RealVis" in self.model_name:
-            prompt = f"RAW photo, interior design, {base_element} {value}, high quality, detailed"
-        else:
-            prompt = f"interior design, {base_element} {value}, high quality, detailed"
-
-        try:
-            return self.pipe(
-                prompt=prompt,
-                negative_prompt="bad quality, blurry, distorted",
-                image=image,
-                strength=strength,
-                num_inference_steps=25,
-                guidance_scale=7.5
-            ).images[0]
-        except Exception as e:
-            print(f"Ошибка change_element: {e}")
-            return image