ControlNet / app.py
trashchenkov's picture
Update app.py
7b4f078 verified
import gradio as gr
import numpy as np
import torch
import re
from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline, ControlNetModel
from peft import PeftModel
# Устройство и тип данных
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
# Изначально загружаем модель по умолчанию (без ControlNet)
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-модификации (unet + text_encoder)
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==True, загружается pipeline с ControlNet.
- В зависимости от controlnet_mode выбирается модель контроля.
- Если use_ip_adapter==True, к pipeline добавляется IP-Adapter с указанным масштабом.
"""
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, создаём новый pipeline с ним
if use_controlnet:
# Выбираем модель ControlNet в зависимости от режима
if controlnet_mode == "edge_detection":
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-canny",
cache_dir="./models_cache",
torch_dtype=torch_dtype
)
elif controlnet_mode == "pose_estimation":
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-openpose",
cache_dir="./models_cache",
torch_dtype=torch_dtype
)
else:
# По умолчанию используем edge_detection
controlnet = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-canny",
cache_dir="./models_cache",
torch_dtype=torch_dtype
)
try:
# Создаем pipeline с ControlNet
pipeline = StableDiffusionControlNetPipeline.from_pretrained(
model,
torch_dtype=torch_dtype,
controlnet=controlnet,
safety_checker=None
).to(device)
except Exception as e:
raise gr.Error(f"Ошибка при создании ControlNet pipeline: {e}")
# Подгружаем LoRA для нового pipeline
try:
pipeline.unet = PeftModel.from_pretrained(pipeline.unet, "./unet")
pipeline.text_encoder = PeftModel.from_pretrained(pipeline.text_encoder, "./text_encoder")
except Exception as e:
print(f"Не удалось подгрузить LoRA в ControlNet pipeline: {e}")
# Если включён IP-Adapter, загружаем его
if use_ip_adapter:
try:
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
pipeline.set_ip_adapter_scale(ip_adapter_scale)
except Exception as e:
raise gr.Error(f"Ошибка при загрузке IP-Adapter: {e}")
# Пытаемся сгенерировать изображение через pipeline с ControlNet
try:
image = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
image=controlnet_image, # изображение, загруженное пользователем для ControlNet
controlnet_conditioning_scale=float(control_strength),
ip_adapter_image=ip_adapter_image if use_ip_adapter else None,
).images[0]
except Exception as e:
raise gr.Error(f"Ошибка при генерации изображения через ControlNet pipeline: {e}")
else:
# Если ControlNet не включён, используем базовый pipeline
pipeline = pipe
# Если включён IP-Adapter, загружаем его в базовый pipeline
if use_ip_adapter:
try:
pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter-plus_sd15.bin")
pipeline.set_ip_adapter_scale(ip_adapter_scale)
except Exception as e:
raise gr.Error(f"Ошибка при загрузке IP-Adapter: {e}")
try:
image = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
num_inference_steps=num_inference_steps,
width=width,
height=height,
generator=generator,
ip_adapter_image=ip_adapter_image if use_ip_adapter else None,
).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 with Control Options")
# Поле для ввода/смены модели
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)
# Продвинутые настройки
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
with gr.Accordion("ControlNet Settings", open=False):
use_controlnet = gr.Checkbox(label="Enable ControlNet", value=False)
control_strength = gr.Slider(label="ControlNet Conditioning Scale", minimum=0.0, maximum=2.0, step=0.1, value=0.9)
controlnet_mode = gr.Dropdown(label="ControlNet Mode", choices=["edge_detection", "pose_estimation"], value="edge_detection")
controlnet_image = gr.Image(label="Upload ControlNet Image", type="pil")
# Дополнительные настройки для IP-Adapter
with gr.Accordion("IP-Adapter Settings", open=False):
use_ip_adapter = gr.Checkbox(label="Enable IP-Adapter", value=False)
ip_adapter_scale = gr.Slider(label="IP-Adapter Scale", minimum=0.0, maximum=2.0, step=0.1, value=0.6)
ip_adapter_image = gr.Image(label="Upload IP-Adapter Image", type="pil")
# Примеры
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()