Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 import random
 import time
-import re
 from diffusers import DiffusionPipeline, LCMScheduler
 
 # -------------------------------
@@ -25,80 +24,60 @@ pipe.enable_vae_slicing()
 pipe.set_progress_bar_config(disable=True)
 
 # -------------------------------
-# PROMPT UNDERSTANDING ENGINE
+# FAST THINKING PROMPT ENGINE
 # -------------------------------
-def refine_prompt(user_prompt: str):
+def refine_prompt_fast(user_prompt: str):
     """
-    Deterministically converts user input to a structured prompt and negative prompt.
-    Ensures SD generates exactly the object the user wants.
+    Instant deterministic prompt refinement (<1ms)
     """
-    # Lowercase and strip
-    p = user_prompt.lower().strip()
-
-    # Attempt to extract a single object from known list
-    known_objects = ["apple","banana","snake","cat","dog","fox","rabbit","dragon","bird","frog","hamster"]
-    obj_match = next((w for w in known_objects if w in p), None)
-    subject = obj_match if obj_match else p
+    # Known objects
+    known_objects = {"apple","banana","snake","cat","dog","fox","rabbit","dragon","bird","frog","hamster"}
+    object_found = next((w for w in known_objects if w in user_prompt.lower()), None)
+    subject = object_found if object_found else user_prompt.lower().strip()
 
     # Style detection
-    is_cute = any(w in p for w in ["cute","adorable","kawaii"])
-    is_realistic = any(w in p for w in ["realistic","photo","photograph"])
-    is_cartoon = any(w in p for w in ["cartoon","anime","illustration"])
-
-    # Base prompt template
-    prompt = f"a single {subject}, centered, isolated"
-
-    if is_cute:
-        prompt += ", cute, friendly, rounded body, big expressive eyes, soft lighting, smooth cartoon style, pastel colors"
-    elif is_cartoon:
-        prompt += ", cartoon style, clean lines, vibrant colors, simple background"
-    elif is_realistic:
-        prompt += ", ultra realistic, natural anatomy, professional photography"
+    style = ""
+    if any(w in user_prompt.lower() for w in ["cute","adorable","kawaii"]):
+        style = ", cute, friendly, rounded body, big eyes, pastel colors, cartoon style"
+    elif any(w in user_prompt.lower() for w in ["realistic","photo","photograph"]):
+        style = ", ultra realistic, natural anatomy, professional photography"
     else:
-        prompt += ", high quality, detailed, clean background"
+        style = ", high quality, clean background"
 
-    # Strong negative prompt to prevent hallucinations
-    negative = (
-        "multiple objects, duplicate, blurry, low quality, cropped, out of frame, "
-        "horror, grotesque, aggressive, scary, weird colors, artifacts"
-    )
+    # Construct prompt
+    prompt = f"a single {subject}, centered, isolated{style}"
+
+    # Negative prompt (always same)
+    negative = "multiple objects, duplicate, blurry, low quality, cropped, out of frame, horror, grotesque, aggressive, weird colors, artifacts"
 
     return prompt, negative
 
 # -------------------------------
-# ETA CALCULATION
+# REALISTIC ETA CALCULATION
 # -------------------------------
 def estimate_time(steps, resolution):
-    # CPU empirical timing (seconds per step)
-    per_step = {
-        256: 6,
-        512: 12,
-        768: 25,
-        1024: 45
-    }[int(resolution)]
-
-    overhead = 10  # initial model load / conditioning
+    per_step = {256:6, 512:12, 768:25, 1024:45}[int(resolution)]
+    overhead = 2  # thinking stage is <1s, plus small overhead
     est = overhead + steps * per_step
     minutes = est // 60
     seconds = est % 60
     return f"⏱️ Estimated time: ~{int(minutes)}m {int(seconds)}s"
 
 # -------------------------------
-# GENERATION FUNCTION (with live status)
+# IMAGE GENERATION WITH LIVE STATUS
 # -------------------------------
 def generate(prompt, resolution, steps):
+    # --- THINKING PHASE ---
     start_time = time.time()
     yield None, "🧠 Understanding your prompt..."
-    
-    refined_prompt, neg_prompt = refine_prompt(prompt)
-    
+    refined_prompt, neg_prompt = refine_prompt_fast(prompt)
     yield None, "🎨 Generating image (CPU, please wait)..."
-    
+
+    # --- GENERATION PHASE ---
     seed = random.randint(0, 10**9)
     gen = torch.Generator("cpu").manual_seed(seed)
-
-    pipe.scheduler.set_timesteps(int(steps))  # ensure LCM fast path
     
+    pipe.scheduler.set_timesteps(int(steps))  # fast path
     img = pipe(
         prompt=refined_prompt,
         negative_prompt=neg_prompt,
@@ -108,7 +87,7 @@ def generate(prompt, resolution, steps):
         height=int(resolution),
         generator=gen
     ).images[0]
-    
+
     duration = int(time.time() - start_time)
     yield [img], f"✅ Finished in {duration}s | Seed: {seed}"