oldnewmodels.py

import torch
from diffusers import (
    StableDiffusionXLImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    DDIMScheduler,
    PNDMScheduler,
    EulerDiscreteScheduler,
    DPMSolverMultistepScheduler
)
from PIL import Image, ImageFilter, ImageEnhance
import numpy as np
import cv2

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', 'T4'])
        
        # Основная модель - RealVisXL V4 
        print(f"Loading {self.model_name} on {gpu_name}...")
        self.pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
            "SG161222/RealVisXL_V4.0",
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16"
        ).to(self.device)
        
        # БЕЗ ЭТИХ СТРОК! Они замедляют H200!
        # self.pipe.enable_model_cpu_offload()
        # self.pipe.enable_vae_slicing()
        
        # Inpainting модель
        try:
            self.inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
                "stabilityai/stable-diffusion-2-inpainting",
                torch_dtype=torch.float16,
                safety_checker=None
            ).to(self.device)
            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["Современный минимализм"])

        # Строка 56-57 должна быть:
        if image.width > 768 or image.height > 768:
        image.thumbnail((768, 768), Image.Resampling.LANCZOS)

        # Оптимальные настройки для H200:

             quality_settings = {
            "fast": {"steps": 15, "guidance": 7.0},
            "balanced": {"steps": 20, "guidance": 8.0},
            "ultra": {"steps": 30, "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
        result = self.pipe(
            prompt=full_prompt,
            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"],
            original_size=(768, 768),
            target_size=(768, 768)
        ).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_2="interior design variation, same style, different details",
                image=image,
                strength=0.4 + (i * 0.05),
                num_inference_steps=20,
                guidance_scale=7.5
            ).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"
        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=30,
            guidance_scale=8.0
        ).images[0]
        
        return result
    
    def create_style_comparison(self, image, styles, quality="fast"):
        """Создание сравнения стилей"""
        results = []
        
        # Настройки для быстрой генерации
        steps = 15 if quality == "fast" else 25
        
        for style in styles:
            styled = self.apply_style_pro(
                image, 
                style, 
                "living room",
                strength=0.75,
                quality=quality
            )
            results.append((style, styled))
            
        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:
    """Класс для удаления объектов"""
    
    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):
        """Генерация маски на основе текстового описания"""
        # Простая маска в центре (заглушка)
        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)
        
        # Рисуем круг
        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