Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import random | |
from diffusers import DiffusionPipeline, LMSDiscreteScheduler | |
import torch | |
import time # Import time for measuring generation time | |
# Device configuration (explicitly set to CPU) | |
DEVICE = "cpu" | |
# Model Options (optimized for CPU) | |
MODEL_OPTIONS = { | |
"Medium Quality (Faster)": "stabilityai/stable-diffusion-2-base", | |
"Fastest (Draft Quality)": "hf-internal-testing/tiny-stable-diffusion-pipe", | |
} | |
# Default to fastest model | |
DEFAULT_MODEL_ID = MODEL_OPTIONS["Fastest (Draft Quality)"] | |
# Cache models to avoid reloading them for each generation | |
PIPELINES = {} | |
def load_pipeline(model_id): | |
if model_id in PIPELINES: | |
return PIPELINES[model_id] | |
else: | |
pipe = DiffusionPipeline.from_pretrained( | |
model_id, torch_dtype=torch.float32 | |
) | |
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config) | |
pipe = pipe.to(DEVICE) | |
PIPELINES[model_id] = pipe | |
return pipe | |
def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice): | |
if not prompt: | |
raise gr.Error("Будь ласка, введіть опис для зображення.") | |
pipe = load_pipeline(MODEL_OPTIONS[model_choice]) | |
generator = torch.Generator(device=DEVICE) | |
if not randomize_seed: | |
generator = generator.manual_seed(seed) | |
start_time = time.time() # Record start time | |
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 | |
end_time = time.time() | |
generation_time = end_time - start_time | |
return images, f"Час генерації: {generation_time:.2f} секунд" # Return images and generation time | |
# ... (Gradio UI remains largely the same, with an added status text output) | |
run_button = gr.Button("Згенерувати") | |
gallery = gr.Gallery(label="Згенеровані зображення") | |
status_text = gr.Textbox(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, status_text], # Output both the gallery and status text | |
) | |