Update app.py

app.py

@@ -3,6 +3,8 @@ from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDisc
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 import gradio as gr
+from tqdm.auto import tqdm
+import psutil
 
 base = "stabilityai/stable-diffusion-xl-base-1.0"
 repo = "ByteDance/SDXL-Lightning"
@@ -12,24 +14,32 @@ ckpt = "sdxl_lightning_4step_unet.safetensors"
 unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cpu")
 unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float32).to("cpu")
+
+# Ensure sampler uses "trailing" timesteps.
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
 def generate_images(prompt, num_inference_steps, guidance_scale, batch_size):
-    images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_size=batch_size).images
+    with tqdm(total=num_inference_steps, desc="Inference Progress") as pbar:
+        images = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_size=batch_size, progress_bar=pbar).images
     return images
 
 # Define Gradio interface
+def get_cpu_info():
+    cpu_name = psutil.cpu_freq().brand
+    memory_available = psutil.virtual_memory().available // 1024 // 1024
+    return f"CPU: {cpu_name}, Memory: {memory_available} MB"
+
 iface = gr.Interface(
     fn=generate_images,
     inputs=[
         gr.Textbox(label="Prompt"),
-        gr.Slider(label="Num Inference Steps", minimum=1, maximum=50, step=1, value=4),
         gr.Slider(label="Guidance Scale", minimum=0, maximum=20, step=0.1, value=0),
-        gr.Slider(label="Batch Size", minimum=1, maximum=8, step=1, value=2),
+        gr.Slider(label="Batch Size", minimum=1, maximum=4, step=1, value=1),
     ],
     outputs=gr.Gallery(label="Generated Images"),
     title="SDXL Lightning 4-Step Inference (CPU)",
     description="Generate images with Stable Diffusion XL Lightning 4-Step model on CPU.",
+    extra_info=get_cpu_info,
 )
 
 iface.launch()