Spaces:
Running
Running
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() | |