import gradio as gr import numpy as np import torch from diffusers import DiffusionPipeline from peft import PeftModel import re # Устройство и тип данных device = "cuda" if torch.cuda.is_available() else "cpu" torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 # Регулярное выражение для проверки корректности модели VALID_REPO_ID_REGEX = re.compile(r"^[a-zA-Z0-9._\-]+/[a-zA-Z0-9._\-]+$") def is_valid_repo_id(repo_id): return bool(VALID_REPO_ID_REGEX.match(repo_id)) and not repo_id.endswith(('-', '.')) # Базовые константы MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 # Изначально загружаем модель по умолчанию model_repo_id = "CompVis/stable-diffusion-v1-4" pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype, safety_checker=None).to(device) # Попробуем подгрузить LoRA-модификации try: pipe.unet = PeftModel.from_pretrained(pipe.unet, "./unet") pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, "./text_encoder") except Exception as e: print(f"Не удалось подгрузить LoRA по умолчанию: {e}") def infer( model, prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, use_controlnet, control_strength, controlnet_mode, controlnet_image, use_ip_adapter, ip_adapter_scale, ip_adapter_image, progress=gr.Progress(track_tqdm=True), ): """ Функция генерации изображения. Параметры use_controlnet, control_strength, controlnet_mode, controlnet_image, use_ip_adapter, ip_adapter_scale, ip_adapter_image — это заглушки для демонстрации UI. """ global model_repo_id, pipe # Если пользователь ввёл другую модель, пробуем её загрузить с нуля if model != model_repo_id: if not is_valid_repo_id(model): raise gr.Error(f"Некорректный идентификатор модели: '{model}'. Проверьте название.") try: new_pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device) # Повторно подгружаем LoRA try: new_pipe.unet = PeftModel.from_pretrained(new_pipe.unet, "./unet") new_pipe.text_encoder = PeftModel.from_pretrained(new_pipe.text_encoder, "./text_encoder") except Exception as e: raise gr.Error(f"Не удалось подгрузить LoRA: {e}") # Обновляем глобальные переменные pipe = new_pipe model_repo_id = model except Exception as e: raise gr.Error(f"Не удалось загрузить модель '{model}'.\nОшибка: {e}") # Создаём генератор случайных чисел для детерминированности generator = torch.Generator(device=device).manual_seed(seed) # --- Здесь должна быть интеграция ControlNet, IP-adapter и т.д. --- # Для демонстрации интерфейса просто вызываем pipe как обычно. # ------------------------------------------------------------------ try: image = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, ).images[0] except Exception as e: raise gr.Error(f"Ошибка при генерации изображения: {e}") return image, seed # Примеры для удобного тестирования examples = [ "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k", "An astronaut riding a green horse", "A delicious ceviche cheesecake slice", ] # Дополнительный CSS для оформления css = """ #col-container { margin: 0 auto; max-width: 640px; } """ # Создаём Gradio-приложение with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown("# Text-to-Image App") # Поле для ввода/смены модели model = gr.Textbox( label="Model", value="CompVis/stable-diffusion-v1-4", # Значение по умолчанию interactive=True ) # Основные поля для Prompt и Negative Prompt prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) negative_prompt = gr.Text( label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=True, ) # Слайдер для выбора seed seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, ) # Слайдеры для guidance_scale и num_inference_steps guidance_scale = gr.Slider( label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=7.0, ) num_inference_steps = gr.Slider( label="Number of inference steps", minimum=1, maximum=50, step=1, value=20, ) # Кнопка запуска run_button = gr.Button("Run", variant="primary") # Поле для отображения результата result = gr.Image(label="Result", show_label=False) # Продвинутые настройки (Accordion) with gr.Accordion("Advanced Settings", open=False): with gr.Row(): width = gr.Slider( label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) height = gr.Slider( label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=512, ) # --- Дополнительные элементы для ControlNet --- use_controlnet = gr.Checkbox(label="Use ControlNet", value=False) with gr.Group(visible=False) as controlnet_group: control_strength = gr.Slider( label="ControlNet Strength", minimum=0.0, maximum=2.0, step=0.1, value=1.0, ) controlnet_mode = gr.Dropdown( label="ControlNet Mode", choices=["edge_detection", "pose_estimation", "depth_estimation"], value="edge_detection", ) controlnet_image = gr.Image( label="ControlNet Image", source="upload", type="pil" ) # Функция для управления видимостью группы ControlNet def update_controlnet_group(use_controlnet): return {"visible": use_controlnet} use_controlnet.change( update_controlnet_group, inputs=[use_controlnet], outputs=[controlnet_group] ) # --- Дополнительные элементы для IP-adapter --- use_ip_adapter = gr.Checkbox(label="Use IP-adapter", value=False) with gr.Group(visible=False) as ip_adapter_group: ip_adapter_scale = gr.Slider( label="IP-adapter Scale", minimum=0.0, maximum=2.0, step=0.1, value=1.0, ) ip_adapter_image = gr.Image( label="IP-adapter Image", source="upload", type="pil" ) # Функция для управления видимостью группы IP-adapter def update_ip_adapter_group(use_ip_adapter): return {"visible": use_ip_adapter} use_ip_adapter.change( update_ip_adapter_group, inputs=[use_ip_adapter], outputs=[ip_adapter_group] ) # Примеры gr.Examples(examples=examples, inputs=[prompt]) # Связка кнопки "Run" с функцией "infer" run_button.click( infer, inputs=[ model, prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, use_controlnet, control_strength, controlnet_mode, controlnet_image, use_ip_adapter, ip_adapter_scale, ip_adapter_image ], outputs=[result, seed], ) # Запуск if __name__ == "__main__": demo.launch()