Update app.py

app.py

@@ -1,116 +1,75 @@
-from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image
 import torch
 import os
-from PIL import Image
-import numpy as np
-import gradio as gr
+import psutil
 import time
 import math
+import gradio as gr
+from PIL import Image
+from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image
 
-# FORCE CPU SETTINGS
-device = torch.device("cpu")
-torch_device = "cpu"
-# CPUs generally require float32 for stability and compatibility
-torch_dtype = torch.float32 
-
-print(f"Running on DEVICE: {device}")
+# --- System Resource Logic ---
+def get_system_usage():
+    cpu = psutil.cpu_percent(interval=None)
+    ram = psutil.virtual_memory().percent
+    # Returns values for the progress bars
+    return cpu / 100, ram / 100, f"CPU: {cpu}% | RAM: {ram}%"
 
-# Load pipelines
-# We remove variant="fp16" because we are using float32 for CPU
-pipe_kwargs = {
-    "torch_dtype": torch_dtype,
-    "safety_checker": None if os.environ.get("SAFETY_CHECKER") != "True" else None,
-    "use_safetensors": True
-}
+# --- Model Loading (Optimized for 16GB) ---
+device = "cpu"
+torch_dtype = torch.float32
 
-i2i_pipe = AutoPipelineForImage2Image.from_pretrained(
-    "stabilityai/sdxl-turbo", **pipe_kwargs
-)
-t2i_pipe = AutoPipelineForText2Image.from_pretrained(
-    "stabilityai/sdxl-turbo", **pipe_kwargs
-)
+pipe_kwargs = {"torch_dtype": torch_dtype, "use_safetensors": True}
+i2i_pipe = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", **pipe_kwargs)
+t2i_pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", **pipe_kwargs)
 
-# OPTIMIZATION FOR 16GB RAM
-# enables slicing the attention computation into steps to save memory
 i2i_pipe.enable_attention_slicing()
 t2i_pipe.enable_attention_slicing()
+i2i_pipe.enable_vae_tiling() 
+t2i_pipe.enable_vae_tiling()
 
-# Moves models to CPU
-i2i_pipe.to("cpu")
-t2i_pipe.to("cpu")
-
-i2i_pipe.set_progress_bar_config(disable=False) # Enabled so you can see it's working
-t2i_pipe.set_progress_bar_config(disable=False)
-
-def resize_crop(image, size=512):
-    image = image.convert("RGB")
-    w, h = image.size
-    image = image.resize((size, int(size * (h / w))), Image.BICUBIC)
-    return image
-
-async def predict(init_image, prompt, strength, steps, seed=1231231):
+def predict(init_image, prompt, strength, steps, seed):
     generator = torch.manual_seed(seed)
-    last_time = time.time()
-    
-    # Ensure steps are sufficient for Turbo
-    if int(steps * strength) < 1:
-        steps = math.ceil(1 / max(0.10, strength))
-
     if init_image is not None:
-        init_image = resize_crop(init_image)
-        results = i2i_pipe(
-            prompt=prompt,
-            image=init_image,
-            generator=generator,
-            num_inference_steps=int(steps),
-            guidance_scale=0.0,
-            strength=strength,
-            width=512,
-            height=512,
-            output_type="pil",
-        )
+        # Simple square crop for better results
+        w, h = init_image.size
+        s = min(w, h)
+        init_image = init_image.crop(((w-s)//2, (h-s)//2, (w+s)//2, (h+s)//2)).resize((512, 512))
+        return i2i_pipe(prompt=prompt, image=init_image, generator=generator, num_inference_steps=int(steps), 
+                        guidance_scale=0.0, strength=strength).images[0]
     else:
-        results = t2i_pipe(
-            prompt=prompt,
-            generator=generator,
-            num_inference_steps=int(steps),
-            guidance_scale=0.0,
-            width=512,
-            height=512,
-            output_type="pil",
-        )
-    
-    print(f"Inference took {time.time() - last_time:.2f} seconds")
-    return results.images[0]
+        return t2i_pipe(prompt=prompt, generator=generator, num_inference_steps=int(steps), 
+                        guidance_scale=0.0, width=512, height=512).images[0]
 
-# --- Gradio UI Section (Keep your existing CSS and Blocks) ---
-css = """#container{ margin: 0 auto; max-width: 80rem; }"""
-with gr.Blocks(css=css) as demo:
-    # ... (Keep your UI layout exactly the same)
-    # Ensure the button calls the predict function
-    # Note: Removed the automatic 'change' triggers to prevent CPU freezing 
-    # while typing. Better to use the Generate button only.
-    
-    init_image_state = gr.State()
+# --- UI Construction ---
+with gr.Blocks(css="#container{ max-width: 60rem; margin: 0 auto; }") as demo:
     with gr.Column(elem_id="container"):
-        gr.Markdown("# SDXL Turbo CPU Edition")
-        with gr.Row():
-            prompt = gr.Textbox(placeholder="Prompt...", scale=5)
-            generate_bt = gr.Button("Generate", scale=1)
+        gr.Markdown("## 🚀 SDXL Turbo CPU + System Monitor")
         
+        # System Usage Dashboard
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(sources=["upload", "webcam"], type="pil")
+                cpu_bar = gr.Label(label="System Status") # Text label for quick reading
             with gr.Column():
-                image_output = gr.Image(type="filepath")
-                
-                with gr.Accordion("Settings", open=True):
-                    strength = gr.Slider(label="Strength", value=0.7, minimum=0.0, maximum=1.0)
-                    steps = gr.Slider(label="Steps (1-4 is best for Turbo)", value=2, minimum=1, maximum=10, step=1)
-                    seed = gr.Slider(label="Seed", minimum=0, maximum=999999, step=1, randomize=True)
+                cpu_plot = gr.Slider(label="CPU Load", interactive=False)
+                ram_plot = gr.Slider(label="RAM Usage", interactive=False)
+
+        with gr.Row():
+            prompt = gr.Textbox(placeholder="A cinematic cat...", label="Prompt", scale=4)
+            generate_bt = gr.Button("Generate", variant="primary", scale=1)
+        
+        with gr.Row():
+            image_input = gr.Image(type="pil", label="Input (i2i)")
+            image_output = gr.Image(label="Result")
+
+        with gr.Accordion("Settings", open=False):
+            strength = gr.Slider(0.0, 1.0, 0.7, label="Strength")
+            steps = gr.Slider(1, 4, 2, step=1, label="Steps")
+            seed = gr.Number(42, label="Seed")
+
+        # Auto-update logic (Runs every 2 seconds)
+        monitor_timer = gr.Timer(2)
+        monitor_timer.tick(get_system_usage, outputs=[cpu_plot, ram_plot, cpu_bar])
 
-        inputs = [image_input, prompt, strength, steps, seed]
-        generate_bt.click(fn=predict, inputs=inputs, outputs=image_output)
+        generate_bt.click(predict, [image_input, prompt, strength, steps, seed], image_output)
 
-demo.queue()
 demo.launch()