Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import random | |
from diffusers import DiffusionPipeline, DDIMScheduler, LMSDiscreteScheduler | |
import torch | |
# Device configuration (explicitly set to CPU) | |
DEVICE = "cpu" | |
# Model Options (for user selection) | |
MODEL_OPTIONS = { | |
"High Quality (Slower)": "stabilityai/stable-diffusion-xl-base-1.0", | |
"Fast (Lower Quality)": "CompVis/stable-diffusion-v1-4", # Smaller, faster model | |
} | |
# Default to faster model | |
DEFAULT_MODEL_ID = MODEL_OPTIONS["Fast (Lower Quality)"] | |
def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice): | |
model_id = MODEL_OPTIONS[model_choice] | |
# Load Model based on user selection | |
pipe = DiffusionPipeline.from_pretrained( | |
model_id, torch_dtype=torch.float32 | |
) | |
# Use LMSDiscreteScheduler for faster generation on CPU with CompVis model | |
if model_choice == "Fast (Lower Quality)": | |
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config) | |
pipe = pipe.to(DEVICE) | |
generator = torch.Generator(device=DEVICE) | |
if not randomize_seed: | |
generator = generator.manual_seed(seed) | |
images = pipe( | |
prompt, | |
negative_prompt=negative_prompt, | |
width=width, | |
height=height, | |
guidance_scale=guidance_scale, | |
num_inference_steps=num_inference_steps, | |
num_images_per_prompt=num_images, | |
generator=generator, | |
).images | |
return images | |
# Gradio Interface | |
with gr.Blocks(title="Генерація зображень за текстом", theme=gr.themes.Soft()) as demo: | |
gr.Markdown( | |
""" | |
## Text-to-Image Generation 🤖🎨 | |
**Створіть дивовижні зображення зі своєї уяви!** | |
Введіть опис, налаштуйте параметри і дозвольте моделі створити для вас витвір мистецтва. | |
""") | |
with gr.Row(): | |
prompt = gr.Textbox(label="Опис", placeholder="Напишіть ваш опис тут...") | |
negative_prompt = gr.Textbox(label="Негативний опис (необов'язково)") | |
with gr.Row(): | |
seed = gr.Number(label="Початкове число", value=0) | |
randomize_seed = gr.Checkbox(label="Випадкове початкове число", value=True) | |
with gr.Row(): # Added this row for model selection | |
model_choice = gr.Radio(label="Виберіть модель", choices=list(MODEL_OPTIONS.keys()), value=DEFAULT_MODEL_ID) | |
with gr.Row(): | |
width = gr.Slider(label="Ширина", minimum=256, maximum=MAX_IMAGE_SIZE, value=512, step=64) | |
height = gr.Slider(label="Висота", minimum=256, maximum=MAX_IMAGE_SIZE, value=512, step=64) | |
with gr.Accordion("Додаткові налаштування", open=False): | |
with gr.Row(): | |
guidance_scale = gr.Slider(label="Рівень відповідності опису", minimum=0.0, maximum=20.0, value=7.5, step=0.1, info="Наскільки точно модель повинна слідувати опису.") | |
num_inference_steps = gr.Slider(label="Кількість кроків", minimum=10, maximum=100, value=50, step=5, info="Більше кроків може покращити якість, але займе більше часу.") | |
num_images = gr.Slider(label="Кількість зображень", minimum=1, maximum=4, value=1, step=1) | |
run_button = gr.Button("Згенерувати") | |
gallery = gr.Gallery(label="Згенеровані зображення") | |
run_button.click( | |
fn=generate_image, | |
inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice], | |
outputs=gallery, | |
) | |
demo.launch(debug=True) | |