import gradio as gr import cv2 import numpy as np from PIL import Image, ImageEnhance import torch from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler import spaces import os # --- 1. БАЗЫ ЗНАНИЙ --- ROOM_PROMPTS = { "Гостиная": "living room interior, sofa area, coffee table, lounge", "Кухня": "kitchen interior, kitchen cabinets, countertop, sink, appliances", "Спальня": "bedroom interior, bed with pillows, headboard, cozy", "Ванная / Санузел": "bathroom interior, toilet, sink, mirror, tiled walls, faucet", "Детская": "kids room interior, playroom, toys, single bed, colorful", "Кабинет / Офис": "home office interior, desk, office chair, computer, shelves", "Прихожая / Коридор": "hallway interior, corridor, entrance, wardrobe, mirror", "Общественное (Кафе/Лобби)": "public space interior, cafe interior, tables and chairs, restaurant" } STYLE_PROMPTS = { "Scandi (Скандинавский)": "scandinavian style, white walls, light oak wood, beige palette, soft sunlight, hygge, natural materials", "Loft (Индустриальный)": "industrial loft style, concrete walls, exposed brick, black metal accents, leather textures, dramatic lighting, cinematic", "Neoclassic (Неоклассика)": "modern classic interior, luxury, wall moldings, velvet furniture, marble floors, gold details, crystal chandelier, elegant", "Minimalism (Минимализм)": "minimalism style, ultra modern, clean lines, monolithic forms, hidden lighting, decluttered, monochromatic grey and white", "Japandi (Джапанди)": "japandi style, wabi-sabi, warm earth tones, raw wood, textured plaster walls, organic shapes, zen atmosphere, soft shadows", "Dark Luxury (Муди)": "dark moody interior, walnut wood, dark grey walls, warm dim lighting, fireplace atmosphere, rich textures, expensive" } # --- 2. ЗАГРУЗКА МОДЕЛИ --- pipe = None def load_pipeline(): global pipe if pipe is None: print("Загрузка модели...") controlnet = ControlNetModel.from_pretrained( "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32 ) pipe = StableDiffusionControlNetPipeline.from_pretrained( "SG161222/Realistic_Vision_V5.1_noVAE", controlnet=controlnet, torch_dtype=torch.float32 ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) try: pipe.to("cuda") except: pass return pipe load_pipeline() # --- 3. ОБРАБОТКА --- def boost_contrast(image, factor=2.0): enhancer = ImageEnhance.Contrast(image) return enhancer.enhance(factor) def process_canny(image): if image is None: return None contrasted_image = boost_contrast(image, factor=2.0) image_np = np.array(contrasted_image) image_canny = cv2.Canny(image_np, 50, 150) image_canny = image_canny[:, :, None] image_canny = np.concatenate([image_canny, image_canny, image_canny], axis=2) return Image.fromarray(image_canny) # --- 4. ГЕНЕРАЦИЯ --- @spaces.GPU(duration=60) def generate(image_path, room_type, style_name, strength, seed): if image_path is None: return None, None, None image = Image.open(image_path).convert("RGB") original_filename = os.path.splitext(os.path.basename(image_path))[0] if max(image.size) > 512: ratio = 512 / max(image.size) new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio)) image = image.resize(new_size, Image.LANCZOS) selected_room = ROOM_PROMPTS[room_type] selected_style = STYLE_PROMPTS[style_name] final_prompt = f"{selected_room}, {selected_style}, architectural photography, masterpiece, 8k" negative_prompt = "blurry, low quality, distorted perspective, bad anatomy, worst quality, watermark, signature" canny_image = process_canny(image) generator = torch.manual_seed(int(seed)) result = pipe( prompt=final_prompt, negative_prompt=negative_prompt, image=canny_image, num_inference_steps=20, generator=generator, controlnet_conditioning_scale=float(strength) ).images[0] output_filename = f"3D_CROSS.{original_filename}.jpg" result.save(output_filename, format="JPEG", quality=95) return result, canny_image, output_filename # --- 5. JS СКРИПТЫ --- # Шеринг share_js = """ async () => { const imgElement = document.querySelector('.result-image img'); if (!imgElement) { alert('Сначала сгенерируйте изображение!'); return; } try { const canvas = document.createElement('canvas'); canvas.width = imgElement.naturalWidth; canvas.height = imgElement.naturalHeight; const ctx = canvas.getContext('2d'); ctx.drawImage(imgElement, 0, 0); canvas.toBlob(async (blob) => { if (!blob) { alert('Ошибка обработки изображения.'); return; } const file = new File([blob], '3D_CROSS_render.jpg', { type: 'image/jpeg' }); if (navigator.share) { try { await navigator.share({ title: '3D CROSS Render', text: 'Смотри, какой дизайн я сделал в 3D CROSS!', files: [file] }); } catch (err) { console.log('Отмена.'); } } else { try { const clipboardItem = new ClipboardItem({'image/png': blob}); await navigator.clipboard.write([clipboardItem]); alert('Изображение скопировано в буфер! (Ctrl+V)'); } catch (err) { alert('Ваш браузер не поддерживает отправку. Скачайте файл вручную.'); } } }, 'image/jpeg', 0.95); } catch (error) { console.error('Ошибка JS:', error); alert('Ошибка. Попробуйте скачать файл.'); } } """ # --- 6. ИНТЕРФЕЙС --- css = """ /* Скрываем кнопки Download и Share */ .gradio-container button[aria-label="Download"], .gradio-container button[aria-label="Share"] { display: none !important; } /* Картинка */ .result-image img { width: 100% !important; height: auto !important; max-height: 600px !important; object-fit: contain !important; } /* Fullscreen fix */ .gradio-image-preview { background-color: rgba(10, 10, 10, 0.98) !important; z-index: 9999 !important; } .gradio-image-preview img { height: 95vh !important; width: auto !important; max-width: 95vw !important; object-fit: contain !important; margin: auto !important; display: block !important; } /* Убираем отступы вокруг приложения, чтобы оно красиво встало в iframe */ .container { max-width: 100% !important; margin: 0 !important; padding: 10px !important; } """ with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo: with gr.Row(): with gr.Column(scale=1): input_image = gr.Image(label="Исходник", type="filepath", height=400) with gr.Accordion("Настройки проекта", open=True): room_dropdown = gr.Dropdown(choices=list(ROOM_PROMPTS.keys()), value="Гостиная", label="📂 Тип помещения") style_dropdown = gr.Dropdown(choices=list(STYLE_PROMPTS.keys()), value="Scandi (Скандинавский)", label="🎨 Стиль") strength_slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="Жесткость") seed_number = gr.Number(value=42, label="Seed", precision=0) run_btn = gr.Button("СГЕНЕРИРОВАТЬ 🚀", variant="primary", size="lg") with gr.Column(scale=2): result_image = gr.Image( label="Результат (Кликни для увеличения)", type="pil", height=600, elem_classes="result-image" ) with gr.Row(): download_btn = gr.DownloadButton("📥 СКАЧАТЬ JPG", variant="secondary") share_btn = gr.Button("🔗 ОТПРАВИТЬ / КОПИРОВАТЬ") with gr.Accordion("Debug", open=False): canny_debug = gr.Image(label="Карта линий", type="pil") # Генерация run_btn.click( fn=generate, inputs=[input_image, room_dropdown, style_dropdown, strength_slider, seed_number], outputs=[result_image, canny_debug, download_btn] ) # Кнопки JS share_btn.click(None, [], [], js=share_js) if __name__ == "__main__": demo.launch()