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") # ТОЛЬКО GPU! 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', 'A10G']) # Основная модель - 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) # БЕЗ enable_model_cpu_offload() и enable_vae_slicing() - они замедляют H200! # Настройка scheduler для качества self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config) # 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 ).to(self.device) print("Inpainting model loaded") except Exception as e: print(f"Warning: Could not load inpainting model: {e}") print("Using img2img as fallback for object removal") self.inpaint_pipe = None @torch.inference_mode() def apply_style_pro(self, image, style_name, room_type, strength=0.75, quality="balanced", custom_prompt=None, custom_negative=None): """Применение стиля к изображению""" from design_styles import DESIGN_STYLES # Ресайз для скорости original_size = image.size if image.width > 768 or image.height > 768: image.thumbnail((768, 768), Image.Resampling.LANCZOS) if style_name == "custom" and custom_prompt: # Кастомный промпт full_prompt = custom_prompt negative = custom_negative or "low quality, blurry" else: # Предустановленный стиль style = DESIGN_STYLES.get(style_name, DESIGN_STYLES["Современный минимализм"]) 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" negative = style.get("negative", "low quality, blurry") # Настройки качества - оптимизированные для H200 quality_settings = { "fast": {"steps": 15, "guidance": 6.0}, "balanced": {"steps": 20, "guidance": 7.0}, "ultra": {"steps": 30, "guidance": 8.0} } settings = quality_settings.get(quality, quality_settings["balanced"]) # Генерация с SDXL result = self.pipe( prompt=full_prompt, prompt_2=full_prompt, # Для SDXL negative_prompt=negative, negative_prompt_2=negative, # Для SDXL image=image, strength=strength, num_inference_steps=settings["steps"], guidance_scale=settings["guidance"], # SDXL параметры - оптимизированные original_size=(768, 768), target_size=(768, 768) ).images[0] # Возвращаем к оригинальному размеру если нужно if result.size != original_size and max(original_size) <= 1024: result = result.resize(original_size, Image.Resampling.LANCZOS) return result def create_variations(self, image, num_variations=4): """Создание вариаций дизайна""" variations = [] base_seed = torch.randint(0, 1000000, (1,)).item() # Ресайз для скорости if image.width > 768 or image.height > 768: image.thumbnail((768, 768), Image.Resampling.LANCZOS) 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=6.0 ).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 # Ресайз для скорости original_size = image.size if image.width > 768 or image.height > 768: image.thumbnail((768, 768), Image.Resampling.LANCZOS) 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, # Для SDXL negative_prompt=negative, negative_prompt_2=negative, image=image, strength=min(strength, 0.8), # Ограничиваем для скорости num_inference_steps=20, # Оптимизировано для H200 guidance_scale=6.0 ).images[0] # Возвращаем к оригинальному размеру if result.size != original_size: result = result.resize(original_size, Image.Resampling.LANCZOS) return result def create_style_comparison(self, image, styles, quality="fast"): """Создание сравнения стилей""" results = [] # Настройки для быстрой генерации steps = 15 if quality == "fast" else 20 for style in styles: styled = self.apply_style_pro( image, style, "living room", # default strength=0.75, 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) try: # Используем 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.95, # Немного меньше для скорости num_inference_steps=25, # Оптимизировано! guidance_scale=5.0 # Меньше для скорости ).images[0] return result except Exception as e: print(f"Inpainting failed: {e}, using OpenCV fallback") return self.simple_inpaint(image, mask) 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) # Рисуем круг 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