Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import random | |
from diffusers import DiffusionPipeline | |
import torch | |
import time | |
import psutil | |
# Get the number of physical CPU cores (excluding hyperthreads) | |
NUM_CPU_CORES = psutil.cpu_count(logical=True) | |
# Cap the number of threads to the available physical cores | |
MAX_THREADS = min(8, NUM_CPU_CORES) | |
# Device and hardware configuration | |
DEVICE = "cpu" | |
# Model Options (optimized for CPU and memory constraints) | |
MODEL_OPTIONS = { | |
"Medium Quality (Faster)": "stabilityai/stable-diffusion-2-base", | |
"Fastest (Draft Quality)": "hf-internal-testing/tiny-stable-diffusion-pipe", | |
} | |
# Default to fastest model and lower image size for limited resources | |
DEFAULT_MODEL_ID = MODEL_OPTIONS["Fastest (Draft Quality)"] | |
DEFAULT_IMAGE_SIZE = 512 # Lower default resolution | |
# Cache models to avoid reloading | |
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.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("Будь ласка, введіть опис для зображення.") | |
torch.set_num_threads(MAX_THREADS) # Set the maximum number of threads | |
pipe = load_pipeline(MODEL_OPTIONS[model_choice]) | |
# Adjust memory usage based on available RAM | |
torch.cuda.empty_cache() | |
generator = torch.Generator(device=DEVICE) | |
if not randomize_seed: | |
generator = generator.manual_seed(seed) | |
start_time = time.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} секунд" | |
# ... (Gradio interface remains the same) | |
generation_time = end_time - start_time | |
return images, f"Час генерації: {generation_time:.2f} секунд" | |
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 | |
) | |