Update app.py

app.py

@@ -6,157 +6,178 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from diffusers import DiffusionPipeline, LCMScheduler
 from PIL import Image, ImageFilter
 
-# -------------------------------
+# ===============================
 # PASSWORD
-# -------------------------------
+# ===============================
 PASSWORD = "CREEPERIMG"
 
-# -------------------------------
-# PROMPT ENHANCER SETUP
-# -------------------------------
-ENHANCER_MODEL = "HuggingFaceTB/SmolLM-135M-Instruct"
-tokenizer_enhancer = AutoTokenizer.from_pretrained(ENHANCER_MODEL)
-model_enhancer = AutoModelForCausalLM.from_pretrained(ENHANCER_MODEL)
-
-def enhance_text(user_prompt, negative=False):
-    """Enhance user prompt or negative prompt using small LLM with strict instructions"""
-    if not user_prompt.strip():
-        return ""
-
-    if negative:
-        instruction = f"""
-You are an assistant for Stable Diffusion / image generation.
-Rewrite the user negative prompt to make it clear what to avoid visually.
-Do NOT add unrelated commentary.
-
-User input: {user_prompt}
-Enhanced negative prompt (visual only):
-"""
-    else:
-        instruction = f"""
-You are an assistant for Stable Diffusion / image generation.
-Your task is to **only rewrite the user prompt** to add details for generating a clear, visual image.
-Do NOT include unrelated topics, commentary, or references to text content.
-Focus solely on visual aspects.
-
-User input: {user_prompt}
-Enhanced prompt (visual only):
-"""
-
-    inputs = tokenizer_enhancer(instruction, return_tensors="pt")
-    outputs = model_enhancer.generate(
+# ===============================
+# TEXT MODEL (PROMPT ENHANCER)
+# ===============================
+TEXT_MODEL_ID = "HuggingFaceTB/SmolLM-135M-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
+text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_ID)
+
+def enhance_prompt(user_prompt: str) -> str:
+    """
+    Enhances the user prompt for image generation.
+    Returns ONLY the enhanced prompt text.
+    """
+    instruction = (
+        "Please enhance this prompt so it is suitable for an image generator "
+        "that requires clear instructions. Analyse the prompt, and output as "
+        "much visual detail as possible about it.\n\n"
+        f"Prompt to enhance: {user_prompt}\n\n"
+        "Enhanced prompt:"
+    )
+
+    inputs = tokenizer(instruction, return_tensors="pt")
+    outputs = text_model.generate(
         **inputs,
-        max_new_tokens=50,
-        temperature=0.7,
-        do_sample=True
+        max_new_tokens=60,
+        temperature=0.6,
+        do_sample=True,
     )
-    return tokenizer_enhancer.decode(outputs[0], skip_special_tokens=True).strip()
 
-# -------------------------------
-# IMAGE MODEL SETUP (CPU)
-# -------------------------------
-IMG_MODEL_ID = "runwayml/stable-diffusion-v1-5"
-IMG_ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # ⛔️ CRITICAL: extract only what comes AFTER "Enhanced prompt:"
+    if "Enhanced prompt:" in decoded:
+        decoded = decoded.split("Enhanced prompt:")[-1]
+
+    return decoded.strip()
+
+# ===============================
+# IMAGE MODEL (CPU)
+# ===============================
+IMG_MODEL = "runwayml/stable-diffusion-v1-5"
+LCM_LORA = "latent-consistency/lcm-lora-sdv1-5"
 
 pipe = DiffusionPipeline.from_pretrained(
-    IMG_MODEL_ID,
+    IMG_MODEL,
     torch_dtype=torch.float32,
     safety_checker=None
 )
+
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights(IMG_ADAPTER_ID)
+pipe.load_lora_weights(LCM_LORA)
 pipe.to("cpu")
+
 pipe.enable_attention_slicing()
 pipe.enable_vae_slicing()
 pipe.set_progress_bar_config(disable=True)
 
-# -------------------------------
-# ETA
-# -------------------------------
-def estimate_time(steps, resolution):
+# ===============================
+# TIME ESTIMATION
+# ===============================
+def estimate_time(steps, res):
+    res = int(res)
     steps = int(steps)
-    resolution = int(resolution)
-    per_step = {512:12, 768:25, 1024:45}[resolution]
-    overhead = 2
-    est = overhead + steps * per_step
-    mins = est // 60
-    secs = est % 60
-    return f"⏱️ Estimated: ~{mins}m {secs}s"
-
-# -------------------------------
-# GENERATE IMAGE WITH BLUR REVEAL
-# -------------------------------
-def generate(password, prompt, neg_prompt, resolution, steps):
+    base = {512: 12, 768: 25, 1024: 45}[res]
+    total = steps * base + 5
+    return f"⏱️ Estimated: ~{total//60}m {total%60}s"
+
+# ===============================
+# GENERATION FUNCTION
+# ===============================
+def generate(password, prompt, negative, resolution, steps):
+    size = int(resolution)
+
+    # Password gate
     if password != PASSWORD:
-        return [Image.new("RGB", (int(resolution), int(resolution)), (255,255,255))], "❌ Wrong password", ""
-    
-    # 1️⃣ Enhance prompts
-    enhanced_prompt = enhance_text(prompt, negative=False)
-    enhanced_negative = enhance_text(neg_prompt, negative=True)
-    
-    # Show enhanced prompt for debugging
-    yield [Image.new("RGB", (int(resolution), int(resolution)), (255,255,255))], "🟡 Generating...", enhanced_prompt
-
-    # 2️⃣ Generate the image
-    seed = random.randint(0, 10**9)
-    gen = torch.Generator("cpu").manual_seed(seed)
+        return (
+            [Image.new("RGB", (size, size), "white")],
+            "❌ Wrong password",
+            ""
+        )
+
+    # Placeholder while thinking
+    yield (
+        [Image.new("RGB", (size, size), "white")],
+        "🧠 Enhancing prompt...",
+        ""
+    )
+
+    enhanced = enhance_prompt(prompt)
+
+    # Show enhanced prompt immediately
+    yield (
+        [Image.new("RGB", (size, size), "white")],
+        "🎨 Generating image...",
+        enhanced
+    )
+
+    seed = random.randint(0, 1_000_000_000)
+    generator = torch.Generator("cpu").manual_seed(seed)
+
     pipe.scheduler.set_timesteps(int(steps))
 
-    img = pipe(
-        prompt=enhanced_prompt,
-        negative_prompt=enhanced_negative,
+    start = time.time()
+    image = pipe(
+        prompt=enhanced,
+        negative_prompt=negative,
         num_inference_steps=int(steps),
         guidance_scale=1.2,
-        width=int(resolution),
-        height=int(resolution),
-        generator=gen
+        width=size,
+        height=size,
+        generator=generator
     ).images[0]
+    elapsed = int(time.time() - start)
 
-    # 3️⃣ Progressive blur reveal
-    max_blur = 20
+    # Blur reveal (UI-only)
     for i in range(10):
-        blur_pct = 100 - i*10
-        blurred = img.filter(ImageFilter.GaussianBlur(radius=max_blur * blur_pct/100))
-        yield [blurred], "🟢 Revealing...", enhanced_prompt
+        blur = image.filter(ImageFilter.GaussianBlur(radius=(10 - i)))
+        yield (
+            [blur],
+            f"👀 Revealing... ({i+1}/10)",
+            enhanced
+        )
         time.sleep(1)
 
-    # 4️⃣ Final image
-    yield [img], f"✅ Done | Seed: {seed}", enhanced_prompt
+    yield (
+        [image],
+        f"✅ Done in {elapsed}s | Seed {seed}",
+        enhanced
+    )
 
-# -------------------------------
-# GRADIO UI
-# -------------------------------
+# ===============================
+# UI
+# ===============================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 👾 CREEPER AI — SMART IMAGE GENERATOR")
+    gr.Markdown("# 👾 Creeper AI — Image Generator")
     gr.Markdown(
-        "1: The higher the resolution & steps, the longer the image takes to make.\n"
-        "2: The more detailed the prompt and negative prompt, the better the result."
+        "1️⃣ The higher the resolution & steps, the longer the image takes.\n\n"
+        "2️⃣ More detailed prompts = better results."
     )
 
     with gr.Row():
         with gr.Column():
-            password_in = gr.Textbox(label="Password", placeholder="Enter password to enable generation")
-            prompt_in = gr.Textbox(label="Prompt")
-            neg_in = gr.Textbox(label="Negative Prompt")
+            password = gr.Textbox(label="Password", type="password")
+            prompt = gr.Textbox(label="Prompt")
+            negative = gr.Textbox(label="Negative Prompt")
 
             resolution = gr.Radio([512, 768, 1024], value=512, label="Resolution")
-            steps = gr.Slider(6,8,value=6,step=1,label="Steps")
+            steps = gr.Slider(6, 8, value=6, step=1, label="Steps")
 
-            eta = gr.Markdown("⏱️ Estimated: ~1m 0s")
-            gen_btn = gr.Button("Generate")
+            eta = gr.Markdown("⏱️ Estimated: ~1m")
+            generate_btn = gr.Button("Generate")
 
             status = gr.Markdown("🟢 Ready")
-            enhanced_box = gr.Textbox(label="Enhanced Prompt (sent to image model)", interactive=False)
+            enhanced_box = gr.Textbox(
+                label="Enhanced Prompt (exact output from text AI)",
+                interactive=False
+            )
 
         with gr.Column():
             gallery = gr.Gallery(columns=1)
 
-    for ctrl in [steps, resolution]:
-        ctrl.change(estimate_time, [steps, resolution], eta)
+    resolution.change(estimate_time, [steps, resolution], eta)
+    steps.change(estimate_time, [steps, resolution], eta)
 
-    gen_btn.click(
+    generate_btn.click(
         generate,
-        inputs=[password_in, prompt_in, neg_in, resolution, steps],
+        inputs=[password, prompt, negative, resolution, steps],
         outputs=[gallery, status, enhanced_box]
     )