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