Update app.py

app.py

@@ -1,48 +1,51 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline, LMSDiscreteScheduler
+from diffusers import DiffusionPipeline
 import torch
-import time  # Import time for measuring generation time
+import time
 
-# Device configuration (explicitly set to CPU)
+# Device and hardware configuration
 DEVICE = "cpu"
+NUM_CPU_CORES = 2
 
-# Model Options (optimized for 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
+# 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 them for each generation
+# 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.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
-        pipe = pipe.to(DEVICE)
+        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(NUM_CPU_CORES)  # Set PyTorch thread count
 
     pipe = load_pipeline(MODEL_OPTIONS[model_choice])
 
+    # Adjust memory usage based on available RAM
+    torch.cuda.empty_cache()  # Not strictly necessary on CPU, but good practice
+
     generator = torch.Generator(device=DEVICE)
     if not randomize_seed:
         generator = generator.manual_seed(seed)
 
-    start_time = time.time()  # Record start time
+    start_time = time.time()
     images = pipe(
         prompt,
         negative_prompt=negative_prompt,
@@ -50,17 +53,15 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
         height=height,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
-        num_images_per_prompt=num_images,
+        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)
+    return images, f"Час генерації: {generation_time:.2f} секунд"
 
     run_button = gr.Button("Згенерувати")
     gallery = gr.Gallery(label="Згенеровані зображення")