Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |
| 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 | |