File size: 2,523 Bytes
6505421
c1bb50f
 
172c734
c1bb50f
172c734
6505421
f55585f
 
c1bb50f
172c734
f55585f
172c734
 
f55585f
c1bb50f
172c734
 
c1bb50f
172c734
 
f55585f
172c734
 
 
 
 
 
 
f55585f
172c734
 
 
f55585f
172c734
 
 
 
 
f55585f
 
 
 
 
172c734
f55585f
 
 
 
 
 
 
 
 
 
 
172c734
 
f55585f
172c734
f55585f
 
172c734
f55585f
 
 
172c734
c1bb50f
 
f55585f
 
172c734
c1bb50f
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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
    )