Update models.py

models.py

@@ -2,6 +2,8 @@ import torch
 from diffusers import (
     StableDiffusionXLImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
+    DDIMScheduler,
+    PNDMScheduler,
     EulerDiscreteScheduler,
     DPMSolverMultistepScheduler
 )
@@ -18,7 +20,7 @@ class InteriorDesignerPro:
         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', 'T4'])
         
-        # Основная модель - RealVisXL V4 для фотореалистичных интерьеров
+        # Основная модель - RealVisXL V4 
         print(f"Loading {self.model_name} on {gpu_name}...")
         self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
             "SG161222/RealVisXL_V4.0",
@@ -27,57 +29,53 @@ class InteriorDesignerPro:
             variant="fp16"
         ).to(self.device)
         
-        # Настройка scheduler для лучшего качества
-        self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
+        # БЕЗ ЭТИХ СТРОК! Они замедляют H200!
+        # self.pipe.enable_model_cpu_offload()
+        # self.pipe.enable_vae_slicing()
         
-        # Inpainting модель для удаления объектов
+        # Inpainting модель
         try:
             self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                 "stabilityai/stable-diffusion-2-inpainting",
                 torch_dtype=torch.float16,
-                safety_checker=None,
-                requires_safety_checker=False,
-                local_files_only=False,
-                resume_download=True
+                safety_checker=None
             ).to(self.device)
-            print("Inpainting model loaded successfully")
-        except Exception as e:
-            print(f"Warning: Could not load inpainting model: {e}")
-            print("Using img2img as fallback for object removal")
+            print("Inpainting model loaded")
+        except:
+            print("Warning: Using fallback for inpainting")
             self.inpaint_pipe = None
-    
+
+    @torch.inference_mode()
     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}
+            "fast": {"steps": 15, "guidance": 7.0},
+            "balanced": {"steps": 25, "guidance": 8.0},
+            "ultra": {"steps": 40, "guidance": 9.0}
         }
         
         settings = quality_settings.get(quality, quality_settings["balanced"])
         
-        # Генерация промпта с учетом комнаты
+        # Промпт для SDXL
         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, photorealistic"
         
-        # Генерация с параметрами для SDXL
+        # Генерация с SDXL
         result = self.pipe(
             prompt=full_prompt,
-            prompt_2=full_prompt,  # SDXL требует второй промпт
-            negative_prompt=style.get("negative", "low quality, blurry, deformed"),
-            negative_prompt_2=style.get("negative", "low quality, blurry, deformed"),
+            prompt_2=full_prompt,
+            negative_prompt=style.get("negative", "low quality, blurry"),
+            negative_prompt_2=style.get("negative", "low quality, blurry"),
             image=image,
             strength=strength,
             num_inference_steps=settings["steps"],
             guidance_scale=settings["guidance"],
-            # SDXL специфичные параметры
-            original_size=(1024, 1024),
-            target_size=(1024, 1024)
+            original_size=(768, 768),
+            target_size=(768, 768)
         ).images[0]
         
         return result
@@ -91,14 +89,12 @@ class InteriorDesignerPro:
             torch.manual_seed(base_seed + i)
             
             var = self.pipe(
-                prompt="interior design variation, same style, different details, photorealistic",
-                prompt_2="interior design variation, same style, different details, photorealistic",
+                prompt="interior design variation, same style, different details",
+                prompt_2="interior design variation, same style, different details",
                 image=image,
                 strength=0.4 + (i * 0.05),
-                num_inference_steps=30,
-                guidance_scale=7.5,
-                original_size=(1024, 1024),
-                target_size=(1024, 1024)
+                num_inference_steps=20,
+                guidance_scale=7.5
             ).images[0]
             
             variations.append(var)
@@ -144,7 +140,7 @@ class InteriorDesignerPro:
         element_info = ROOM_ELEMENTS.get(element, {})
         prompt_add = element_info.get("prompt_add", element.lower())
         
-        prompt = f"interior with {value} {prompt_add}, professional photo, photorealistic"
+        prompt = f"interior with {value} {prompt_add}, professional photo"
         negative = f"old {element}, damaged, ugly"
         
         result = self.pipe(
@@ -154,10 +150,8 @@ class InteriorDesignerPro:
             negative_prompt_2=negative,
             image=image,
             strength=strength,
-            num_inference_steps=40,
-            guidance_scale=8.0,
-            original_size=(1024, 1024),
-            target_size=(1024, 1024)
+            num_inference_steps=30,
+            guidance_scale=8.0
         ).images[0]
         
         return result
@@ -167,13 +161,13 @@ class InteriorDesignerPro:
         results = []
         
         # Настройки для быстрой генерации
-        steps = 20 if quality == "fast" else 35
+        steps = 15 if quality == "fast" else 25
         
         for style in styles:
             styled = self.apply_style_pro(
                 image, 
                 style, 
-                "living room",  # default
+                "living room",
                 strength=0.75,
                 quality=quality
             )
@@ -182,6 +176,41 @@ class InteriorDesignerPro:
         return results
 
 
+# Динамическое добавление метода для сетки
+def _create_comparison_grid(self, results):
+    """Создание сетки из результатов"""
+    if not results:
+        return None
+        
+    images = [img for _, img in results]
+    titles = [title for title, _ in results]
+    
+    # Определяем размер сетки
+    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
+
+# Добавляем метод к классу
+InteriorDesignerPro._create_comparison_grid = _create_comparison_grid
+
+
 class ObjectRemover:
     """Класс для удаления объектов"""
     
@@ -221,7 +250,6 @@ class ObjectRemover:
     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)
         
@@ -236,59 +264,3 @@ class ObjectRemover:
                       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