Update models.py

models.py

@@ -1,10 +1,9 @@
 import torch
 from diffusers import (
-    StableDiffusionImg2ImgPipeline,
+    StableDiffusionXLImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
-    DDIMScheduler,
-    PNDMScheduler,
-    EulerDiscreteScheduler
+    EulerDiscreteScheduler,
+    DPMSolverMultistepScheduler
 )
 from PIL import Image, ImageFilter, ImageEnhance
 import numpy as np
@@ -14,24 +13,27 @@ class InteriorDesignerPro:
     def __init__(self):
         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(gpu in gpu_name for gpu in ['A100', 'H100', 'RTX 4090', 'RTX 3090'])
-
-        # Основная модель - RealVis V4 для фотореалистичных интерьеров
+        self.is_powerful_gpu = any(gpu in gpu_name for gpu in ['A100', 'H100', 'RTX 4090', 'RTX 3090', 'T4'])
+        
+        # Основная модель - RealVisXL V4 для фотореалистичных интерьеров
         print(f"Loading {self.model_name} on {gpu_name}...")
-        self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
             "SG161222/RealVisXL_V4.0",
             torch_dtype=torch.float16,
-            safety_checker=None,
-            requires_safety_checker=False,
-            local_files_only=False
+            use_safetensors=True,
+            variant="fp16"
         ).to(self.device)
-
+        
+        # Включаем оптимизации для SDXL
+        self.pipe.enable_model_cpu_offload()
+        self.pipe.enable_vae_slicing()
+        
         # Настройка scheduler для лучшего качества
         self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
-
+        
         # Inpainting модель для удаления объектов
         try:
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
@@ -47,118 +49,130 @@ class InteriorDesignerPro:
             print(f"Warning: Could not load inpainting model: {e}")
             print("Using img2img as fallback for object removal")
             self.inpaint_pipe = None
-            
+    
     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 = 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"
-
-        # Генерация
+        full_prompt = f"{style['prompt']}, {room_specific}, {room_type} interior design, professional photo, high quality, 8k, photorealistic"
+        
+        # Генерация с параметрами для SDXL
         result = self.pipe(
             prompt=full_prompt,
-            negative_prompt=style.get("negative", "low quality, blurry"),
+            prompt_2=full_prompt,  # SDXL требует второй промпт
+            negative_prompt=style.get("negative", "low quality, blurry, deformed"),
+            negative_prompt_2=style.get("negative", "low quality, blurry, deformed"),
             image=image,
             strength=strength,
             num_inference_steps=settings["steps"],
-            guidance_scale=settings["guidance"]
+            guidance_scale=settings["guidance"],
+            # SDXL специфичные параметры
+            original_size=(1024, 1024),
+            target_size=(1024, 1024)
         ).images[0]
-
+        
         return result
-
+        
     def create_variations(self, image, num_variations=4):
         """Создание вариаций дизайна"""
         variations = []
         base_seed = torch.randint(0, 1000000, (1,)).item()
-
+        
         for i in range(num_variations):
             torch.manual_seed(base_seed + i)
-
+            
             var = self.pipe(
-                prompt="interior design variation, same style, different details",
+                prompt="interior design variation, same style, different details, photorealistic",
+                prompt_2="interior design variation, same style, different details, photorealistic",
                 image=image,
                 strength=0.4 + (i * 0.05),
                 num_inference_steps=30,
-                guidance_scale=7.5
+                guidance_scale=7.5,
+                original_size=(1024, 1024),
+                target_size=(1024, 1024)
             ).images[0]
-
+            
             variations.append(var)
-
+            
         return variations
-        
+    
     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)
         sharp = enhancer.enhance(1.5)
-
+        
         # Немного увеличиваем контраст
         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"
+        
+        prompt = f"interior with {value} {prompt_add}, professional photo, photorealistic"
         negative = f"old {element}, damaged, ugly"
-
+        
         result = self.pipe(
             prompt=prompt,
+            prompt_2=prompt,
             negative_prompt=negative,
+            negative_prompt_2=negative,
             image=image,
             strength=strength,
             num_inference_steps=40,
-            guidance_scale=8.0
+            guidance_scale=8.0,
+            original_size=(1024, 1024),
+            target_size=(1024, 1024)
         ).images[0]
-
-        return result
         
+        return result
+    
     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, 
@@ -168,23 +182,23 @@ class InteriorDesignerPro:
                 quality=quality
             )
             results.append((style, styled))
-
+            
         return results
 
 
 class ObjectRemover:
     """Класс для удаления объектов"""
-
+    
     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",
@@ -195,34 +209,91 @@ class ObjectRemover:
             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
         width, height = image.size
         mask = Image.new('L', (width, height), 0)
-
+        
         # Создаем маску в центре
         center_x, center_y = width // 2, height // 2
         radius = int(min(width, height) * precision)
-
+        
         # Рисуем круг
-        import ImageDraw
+        from PIL import ImageDraw
         draw = ImageDraw.Draw(mask)
         draw.ellipse([center_x - radius, center_y - radius, 
                       center_x + radius, center_y + radius], fill=255)
-
+        
         return mask
+
+
+# Добавляем метод _create_comparison_grid к классу при импорте
+def _create_comparison_grid(self, images_with_labels):
+    """Создает сетку из изображений с подписями"""
+    if not images_with_labels:
+        return None
+        
+    images = [img for _, img in images_with_labels]
+    labels = [label for label, _ in images_with_labels]
+    
+    # Определяем размер сетки
+    n = len(images)
+    cols = min(3, n)  # Максимум 3 колонки
+    rows = (n + cols - 1) // cols
+    
+    # Размер одного изображения
+    img_width, img_height = images[0].size
+    padding = 20
+    label_height = 40
+    
+    # Создаем холст
+    grid_width = cols * img_width + (cols + 1) * padding
+    grid_height = rows * (img_height + label_height) + (rows + 1) * padding
+    
+    grid = Image.new('RGB', (grid_width, grid_height), 'white')
+    
+    # Добавляем изображения
+    from PIL import ImageDraw, ImageFont
+    draw = ImageDraw.Draw(grid)
+    
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
+    except:
+        font = None
+    
+    for idx, (img, label) in enumerate(zip(images, labels)):
+        row = idx // cols
+        col = idx % cols
+        
+        x = col * img_width + (col + 1) * padding
+        y = row * (img_height + label_height) + (row + 1) * padding
+        
+        # Вставляем изображение
+        grid.paste(img, (x, y))
+        
+        # Добавляем подпись
+        text_x = x + img_width // 2
+        text_y = y + img_height + 5
+        
+        draw.text((text_x, text_y), label, fill='black', font=font, anchor='mt')
+    
+    return grid
+
+# Патчим класс
+InteriorDesignerPro._create_comparison_grid = _create_comparison_grid
+models.py