ImgGen / app.py
Rooni's picture
Update app.py
261c072
import gradio as gr
from diffusers import StableDiffusionPipeline
import torch
# Словарь с путями к моделям
model_paths = {
"Vector Art": "stablediffusionapi/vector-art",
# Добавьте другие модели здесь
}
# Инициализация и загрузка моделей
models = {}
for name, path in model_paths.items():
pipeline = StableDiffusionPipeline.from_pretrained(path)
pipeline.to("cpu")
models[name] = pipeline
# Функция для генерации изображения
def generate_image(prompt, model_name, sampling_steps, cfg, sampling_method, seed, width, height):
# Выбор модели из словаря
model = models[model_name]
# Установка параметров генерации
generator = torch.manual_seed(seed) if seed != -1 else None
# Генерация изображения
image = model(prompt, num_inference_steps=sampling_steps, guidance_scale=cfg, generator=generator).images[0]
# Преобразование изображения для отображения в Gradio
return image
# Создание интерфейса с использованием gr.Blocks
with gr.Blocks() as app:
gr.Markdown("Генератор изображений")
with gr.Tab("Основные настройки"):
with gr.Row():
prompt_input = gr.Textbox(label="Введите текстовый запрос")
model_select = gr.Radio(choices=list(model_paths.keys()), label="Выберите модель", value="Vector Art")
with gr.Tab("Расширенные настройки"):
with gr.Row():
sampling_steps_input = gr.Slider(minimum=1, maximum=30, step=1, label="Sampling Steps", value=20)
cfg_input = gr.Slider(minimum=1, maximum=20, step=1, label="Classifier Free Guidance (CFG)", value=7)
sampling_method_input = gr.Radio(choices=["DDIM", "DDPM", "PLMS"], label="Sampling Method", value="DDIM")
seed_input = gr.Slider(minimum=-1, maximum=100000, step=1, label="Seed", value=-1)
width_input = gr.Slider(minimum=15, maximum=1023, step=1, label="Width", value=512)
height_input = gr.Slider(minimum=15, maximum=1023, step=1, label="Height", value=512)
with gr.Row():
generate_button = gr.Button("Сгенерировать")
output_image = gr.Image(label="Сгенерированное изображение")
# Подключение функции генерации к кнопке
generate_button.click(
generate_image,
inputs=[prompt_input, model_select, sampling_steps_input, cfg_input, sampling_method_input, seed_input, width_input, height_input],
outputs=output_image
)
app.launch()