Update app.py

app.py

@@ -1,121 +1,164 @@
 import gradio as gr
 import torch
 import random
-from diffusers import DiffusionPipeline, LCMScheduler
 import time
 import re
+import numpy as np
+from diffusers import DiffusionPipeline, LCMScheduler
+from PIL import Image
 
-# Load Model
+# -------------------------------------------------
+# MODEL SETUP (CPU-ONLY, FAST)
+# -------------------------------------------------
 model_id = "runwayml/stable-diffusion-v1-5"
 adapter_id = "latent-consistency/lcm-lora-sdv1-5"
 
-pipe = DiffusionPipeline.from_pretrained(model_id, safety_checker=None, torch_dtype=torch.float32)
+pipe = DiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float32,
+    safety_checker=None
+)
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
 pipe.load_lora_weights(adapter_id)
 pipe.to("cpu")
 
-def analyze_prompt_logic(prompt):
-    """
-    Fast Analysis Engine (Under 1s)
-    Determines if the user wants a single object or a complex scene.
-    """
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+
+# -------------------------------------------------
+# FAST PROMPT ANALYZER (NO DELAYS)
+# -------------------------------------------------
+def analyze_prompt(prompt: str):
     p = prompt.lower()
-    # Check for number intent
-    has_count = re.search(r'\b(one|1|single|sole)\b', p)
-    is_scene = any(k in p for k in ["room", "forest", "city", "street", "landscape", "background", "detailed"])
-    
-    if has_count and not is_scene:
-        # Inject strict focus for single objects
-        return f"((centered single isolated {prompt})), minimalist background, sharp focus, masterpiece"
-    return f"{prompt}, high quality, highly detailed"
-
-def generate_process(prompt, neg_prompt, size, steps, is_random, manual_seed, bulk_count, auto_res):
-    all_images = []
-    start_total = time.time()
-    
-    # --- ANALYSIS PHASE (The 'Thinking' part) ---
-    yield None, "🤔 **Thinking... Analyzing prompt intent**"
-    time.sleep(1.5) # Simulated for UI feedback, actual logic is instant
-    
-    optimized_prompt = analyze_prompt_logic(prompt)
-    
-    # Auto-Resolution Logic
-    if auto_res:
-        w, h = (448, 448) if "scene" not in optimized_prompt else (512, 512)
+
+    # Detect single-object intent
+    single_object = not re.search(r"\b(two|three|multiple|many|group|pile)\b", p)
+
+    # Detect scene
+    is_scene = any(k in p for k in [
+        "room", "street", "city", "forest", "landscape", "background"
+    ])
+
+    if single_object and not is_scene:
+        optimized = (
+            f"single centered {prompt}, isolated object, plain background, "
+            f"sharp focus, product photo"
+        )
+        neg = (
+            "multiple objects, duplicates, crowd, background objects, clutter"
+        )
     else:
-        w, h = int(size), int(size)
-
-    # --- GENERATION PHASE ---
-    for i in range(int(bulk_count)):
-        yield all_images, f"🧠 **Thinking... Strategy: {optimized_prompt[:40]}...**"
-        
-        current_seed = random.randint(0, 10**6) if (is_random or bulk_count > 1) else int(manual_seed)
-        gen = torch.Generator(device="cpu").manual_seed(current_seed)
-        
-        img = pipe(
-            prompt=optimized_prompt, 
-            negative_prompt=neg_prompt if neg_prompt else "duplicate, multiple, split, blurry",
-            num_inference_steps=int(steps),
-            guidance_scale=2.5, # High guidance for strict instruction following
-            width=w, height=h,
-            generator=gen
-        ).images[0]
-        
-        all_images.append(img)
-        yield all_images, f"🎨 **Generating... ({i+1}/{bulk_count})**"
-    
-    total_duration = round(time.time() - start_total, 2)
-    yield all_images, f"✅ **Finished in {total_duration}s** | Seed: {current_seed}"
-
-def update_estimate(steps, size, bulk, auto):
-    # Precise 2-vCPU calibration
-    res_val = 448 if auto else int(size)
-    time_per_step = 14.8 * ((res_val / 512) ** 2)
-    # Add 2 seconds for the Analysis Phase
-    est = round((steps * time_per_step * bulk) + 2, 1)
-    return f"⚡ **Estimated Wait:** ~{est} seconds"
-
-css = """
-.gradio-container { background-color: #0b0f19 !important; color: #00ffcc !important; }
-.generate-btn { background: linear-gradient(90deg, #00ffcc, #0077ff) !important; color: white !important; font-weight: bold !important; border: none !important; }
-.side-panel { border: 1px solid #00ffcc; padding: 15px; border-radius: 12px; background: rgba(0,255,204,0.05); }
-#title { text-shadow: 0 0 10px #00ffcc; }
-"""
-
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
-    gr.Markdown("# 👾 CREEPER AI - V2.3 LOGIC CORE", elem_id="title")
-    
+        optimized = f"{prompt}, high quality"
+        neg = "blurry, low quality"
+
+    return optimized, neg, is_scene
+
+# -------------------------------------------------
+# SMART AUTO-CROP (FAST, NO ML)
+# -------------------------------------------------
+def smart_crop(img: Image.Image, padding=5):
+    gray = np.array(img.convert("L"))
+    mask = gray > 20  # detect non-black-ish pixels
+
+    if not mask.any():
+        return img  # safety fallback
+
+    coords = np.argwhere(mask)
+    y0, x0 = coords.min(axis=0)
+    y1, x1 = coords.max(axis=0)
+
+    h, w = gray.shape
+    x0 = max(0, x0 - padding)
+    y0 = max(0, y0 - padding)
+    x1 = min(w, x1 + padding)
+    y1 = min(h, y1 + padding)
+
+    return img.crop((x0, y0, x1, y1))
+
+# -------------------------------------------------
+# GENERATION CORE (FAST <10s)
+# -------------------------------------------------
+def generate(
+    prompt,
+    size,
+    steps,
+    random_seed,
+    seed_val,
+    auto_res
+):
+    start = time.time()
+
+    optimized_prompt, neg_prompt, is_scene = analyze_prompt(prompt)
+
+    # Resolution logic
+    if auto_res and not is_scene:
+        width = height = 448
+    else:
+        width = height = int(size)
+
+    seed = random.randint(0, 999999) if random_seed else int(seed_val)
+    gen = torch.Generator("cpu").manual_seed(seed)
+
+    image = pipe(
+        prompt=optimized_prompt,
+        negative_prompt=neg_prompt,
+        num_inference_steps=int(steps),
+        guidance_scale=1.2,  # LCM sweet spot
+        width=width,
+        height=height,
+        generator=gen
+    ).images[0]
+
+    # Auto-crop ONLY for single objects
+    if auto_res and not is_scene:
+        image = smart_crop(image, padding=5)
+
+    duration = round(time.time() - start, 2)
+    status = f"✅ Done in {duration}s | Seed: {seed}"
+
+    return [image], status
+
+# -------------------------------------------------
+# ETA (CPU-REALISTIC)
+# -------------------------------------------------
+def estimate_time(steps, size, auto):
+    base = 0.9  # model overhead
+    res = 448 if auto else int(size)
+    step_cost = 0.22 * (res / 512) ** 2
+    est = base + steps * step_cost
+    return f"⚡ Estimated: ~{round(est, 1)}s"
+
+# -------------------------------------------------
+# UI
+# -------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 👾 CREEPER AI — FAST SMART IMAGE CORE")
+
     with gr.Row():
-        with gr.Column(elem_classes="side-panel"):
-            prompt_in = gr.Textbox(label="Main Instruction", placeholder="e.g. One red apple", lines=3)
-            neg_in = gr.Textbox(label="Negative Prompt", placeholder="Optional...")
-            
-            with gr.Accordion("⚙️ Engine Parameters", open=True):
-                bulk_slider = gr.Slider(1, 3, value=1, step=1, label="Bulk Count")
-                auto_res_check = gr.Checkbox(label="Auto-Optimize Resolution", value=False)
-                size_choice = gr.Radio(choices=[512, 768], label="Fixed Resolution", value=512)
-                step_slider = gr.Slider(1, 6, value=4, step=1, label="Steps")
-                
-                with gr.Row():
-                    random_seed = gr.Checkbox(label="Random Seed", value=True)
-                    seed_val = gr.Number(label="Manual Seed", visible=False, value=42)
-            
-            est_box = gr.Markdown("⚡ **Estimated Wait:** ~61.2 seconds")
-            gen_btn = gr.Button("INITIALIZE GENERATION", variant="primary", elem_classes="generate-btn")
+        with gr.Column():
+            prompt_in = gr.Textbox(label="Prompt", placeholder="apple")
+            auto_res = gr.Checkbox(label="Smart Auto-Resolution", value=True)
+            size = gr.Radio([512, 768], value=512, label="Base Resolution")
+            steps = gr.Slider(2, 6, value=4, step=1, label="Steps")
+
+            random_seed = gr.Checkbox(label="Random Seed", value=True)
+            seed_val = gr.Number(label="Manual Seed", value=42, visible=False)
+
+            eta = gr.Markdown("⚡ Estimated: ~2.0s")
+            btn = gr.Button("Generate")
 
         with gr.Column():
-            status = gr.Markdown("🟢 **System Ready**")
-            gallery = gr.Gallery(label="Output", columns=1, height="auto", preview=True)
+            status = gr.Markdown("🟢 Ready")
+            gallery = gr.Gallery(columns=1)
 
     random_seed.change(lambda x: gr.update(visible=not x), random_seed, seed_val)
-    
-    for ctrl in [bulk_slider, auto_res_check, size_choice, step_slider]:
-        ctrl.change(update_estimate, [step_slider, size_choice, bulk_slider, auto_res_check], est_box)
+    for ctrl in [steps, size, auto_res]:
+        ctrl.change(estimate_time, [steps, size, auto_res], eta)
 
-    gen_btn.click(
-        fn=generate_process,
-        inputs=[prompt_in, neg_in, size_choice, step_slider, random_seed, seed_val, bulk_slider, auto_res_check],
+    btn.click(
+        generate,
+        inputs=[prompt_in, size, steps, random_seed, seed_val, auto_res],
         outputs=[gallery, status]
     )
 
-demo.launch()
+demo.launch()