Update app.py

app.py

@@ -15,11 +15,12 @@ def get_torch_dtype():
 def get_device():
     return "cuda" if torch.cuda.is_available() else "cpu"
 
-# === Load the diffusion model ===
-torch_dtype = get_torch_dtype()
-device = get_device()
-pipe = DiffusionPipeline.from_pretrained(MODEL_REPO_ID, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+# === Lazy load the diffusion model ===
+def get_pipe():
+    if not hasattr(get_pipe, "pipe"):
+        pipe = DiffusionPipeline.from_pretrained(MODEL_REPO_ID, torch_dtype=get_torch_dtype()).to(get_device())
+        get_pipe.pipe = pipe
+    return get_pipe.pipe
 
 # === Define custom prompt builder ===
 def build_prompt(word):
@@ -35,15 +36,16 @@ def build_prompt(word):
 # === Image generation function ===
 def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, width, height, seed):
     generator = torch.Generator().manual_seed(seed)
-    return pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    with torch.inference_mode():
+        return get_pipe()(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
 
 # === Inference wrapper ===
 def infer(
@@ -103,7 +105,7 @@ with gr.Blocks(css=css) as demo:
 
             with gr.Row():
                 guidance_scale = gr.Slider(label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=3.5)
-                num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=20)
+                num_inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=4)
 
     run_button.click(
         fn=infer,