Update app.py

app.py

@@ -1,41 +1,48 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline, DDIMScheduler, LMSDiscreteScheduler
+from diffusers import DiffusionPipeline, LMSDiscreteScheduler
 import torch
+import time  # Import time for measuring generation time
 
 # Device configuration (explicitly set to CPU)
 DEVICE = "cpu"
-# Maximum Image Size (Defined at the top level)
-MAX_IMAGE_SIZE = 1024
 
-# Model Options (for user selection)
+# Model Options (optimized for CPU)
 MODEL_OPTIONS = {
-    "High Quality (Slower)": "stabilityai/stable-diffusion-xl-base-1.0",
-    "Fast (Lower Quality)": "CompVis/stable-diffusion-v1-4",  # Smaller, faster model
+    "Medium Quality (Faster)": "stabilityai/stable-diffusion-2-base",
+    "Fastest (Draft Quality)": "hf-internal-testing/tiny-stable-diffusion-pipe",
 }
 
-# Default to faster model
-DEFAULT_MODEL_ID = MODEL_OPTIONS["Fast (Lower Quality)"]
+# Default to fastest model
+DEFAULT_MODEL_ID = MODEL_OPTIONS["Fastest (Draft 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
-    )
+# Cache models to avoid reloading them for each generation
+PIPELINES = {}
 
-    # Use LMSDiscreteScheduler for faster generation on CPU with CompVis model
-    if model_choice == "Fast (Lower Quality)":
+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
 
-    pipe = pipe.to(DEVICE)
+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,
@@ -47,45 +54,20 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
         generator=generator,
     ).images
 
-    return images
-
-# Gradio Interface
-with gr.Blocks(title="Генерація зображень за текстом", theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        ## Text-to-Image Generation 🤖🎨
-        **Створіть дивовижні зображення зі своєї уяви!**  
-        Введіть опис, налаштуйте параметри і дозвольте моделі створити для вас витвір мистецтва.
-        """)
+    end_time = time.time()
+    generation_time = end_time - start_time
 
-    with gr.Row():
-        prompt = gr.Textbox(label="Опис", placeholder="Напишіть ваш опис тут...")
-        negative_prompt = gr.Textbox(label="Негативний опис (необов'язково)")
+    return images, f"Час генерації: {generation_time:.2f} секунд"  # Return images and generation time
 
-    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)
+# ... (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,
+        outputs=[gallery, status_text],  # Output both the gallery and status text
     )
-
-demo.launch(debug=True)