Update app.py

app.py

@@ -0,0 +1,116 @@
+from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image
+import torch
+import os
+from PIL import Image
+import numpy as np
+import gradio as gr
+import time
+import math
+
+# 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}")
+
+# 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
+}
+
+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()
+
+# 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):
+    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",
+        )
+    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]
+
+# --- 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()
+    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)
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(sources=["upload", "webcam"], type="pil")
+            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)
+
+        inputs = [image_input, prompt, strength, steps, seed]
+        generate_bt.click(fn=predict, inputs=inputs, outputs=image_output)
+
+demo.queue()
+demo.launch()