Upload models.py

models.py

@@ -0,0 +1,384 @@
+import torch
+from diffusers import (
+    StableDiffusionXLImg2ImgPipeline,
+    StableDiffusionImg2ImgPipeline,
+    AutoPipelineForImage2Image,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    AutoencoderKL
+)
+from PIL import Image
+import numpy as np
+
+class InteriorDesignerPro:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "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
+        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:
+            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 загружен")
+            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, "
+        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
+
+        # Генерация
+        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
+        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)
+
+        return variations
+
+    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