Update app.py

app.py

@@ -2,113 +2,120 @@ import gradio as gr
 import torch
 import random
 import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from diffusers import DiffusionPipeline, LCMScheduler
 from PIL import Image, ImageFilter
 
 # -------------------------------
-# MODEL SETUP (CPU SAFE)
+# SMALL PROMPT ENHANCER (CPU)
 # -------------------------------
-MODEL_ID = "runwayml/stable-diffusion-v1-5"
-ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"
+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, prefix):
+    """
+    Uses a tiny LLM to rewrite user input into a more detailed instruction.
+    """
+    instruction = f"Rewrite this for an image generator with detail: {user_prompt}"
+    if prefix:
+        instruction = f"{prefix} {user_prompt}"
+
+    inputs = tokenizer_enhancer(instruction, return_tensors="pt")
+    outputs = model_enhancer.generate(
+        **inputs,
+        max_new_tokens=50,
+        temperature=0.7,
+        do_sample=True
+    )
+    text = tokenizer_enhancer.decode(outputs[0], skip_special_tokens=True)
+    return text.strip()
+
+# -------------------------------
+# IMAGE MODEL SETUP (CPU SAFE)
+# -------------------------------
+IMG_MODEL_ID = "runwayml/stable-diffusion-v1-5"
+IMG_ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"
 
 pipe = DiffusionPipeline.from_pretrained(
-    MODEL_ID,
+    IMG_MODEL_ID,
     torch_dtype=torch.float32,
     safety_checker=None
 )
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights(ADAPTER_ID)
+pipe.load_lora_weights(IMG_ADAPTER_ID)
 pipe.to("cpu")
-
 pipe.enable_attention_slicing()
 pipe.enable_vae_slicing()
 pipe.set_progress_bar_config(disable=True)
 
 # -------------------------------
-# PROMPT REFINEMENT
-# -------------------------------
-def refine_prompt(user_prompt: str):
-    """
-    CPU-fast prompt enhancement.
-    Always focuses on a single centered object.
-    """
-    subject = user_prompt.strip()
-    prompt = f"a single {subject}, centered, isolated, high quality, clean background"
-    negative = "multiple objects, duplicate, blurry, low quality, cropped, out of frame, horror, grotesque, weird colors, artifacts"
-    return prompt, negative
-
-# -------------------------------
-# ETA ESTIMATION
+# ETA
 # -------------------------------
 def estimate_time(steps, resolution):
     steps = int(steps)
     resolution = int(resolution)
-    per_step = {256:6, 512:12, 768:25, 1024:45}[resolution]
+    per_step = {512:12, 768:25, 1024:45}[resolution]
     overhead = 2
     est = overhead + steps * per_step
-    minutes = est // 60
-    seconds = est % 60
-    return f"⏱️ Estimated time: ~{int(minutes)}m {int(seconds)}s"
+    mins = est // 60
+    secs = est % 60
+    return f"⏱️ Estimated: ~{mins}m {secs}s"
 
 # -------------------------------
-# IMAGE GENERATION WITH PROGRESSIVE BLUR
+# GENERATE WITH PROGRESSIVE BLUR
 # -------------------------------
-def generate(prompt, resolution, steps):
-    steps = int(steps)
-    resolution = int(resolution)
-
-    # 1️⃣ Refine prompt
-    refined_prompt, neg_prompt = refine_prompt(prompt)
+def generate(prompt, neg_prompt, resolution, steps):
+    # 1️⃣ AI PROMPT ENHANCEMENT
+    enhanced_prompt = enhance_text(prompt, "")
+    enhanced_negative = enhance_text(neg_prompt, "Rewrite negative prompt:")
 
-    # 2️⃣ White placeholder image
-    placeholder = Image.new("RGB", (resolution, resolution), (255, 255, 255))
-    start_time = time.time()
+    # 2️⃣ White placeholder
+    placeholder = Image.new("RGB", (int(resolution), int(resolution)), (255,255,255))
     yield [placeholder], "🟡 Generating..."
 
-    # 3️⃣ CPU Image generation
+    # 3️⃣ CPU IMAGE GENERATION
     seed = random.randint(0, 10**9)
     gen = torch.Generator("cpu").manual_seed(seed)
-    pipe.scheduler.set_timesteps(steps)
+    pipe.scheduler.set_timesteps(int(steps))
 
     img = pipe(
-        prompt=refined_prompt,
-        negative_prompt=neg_prompt,
-        num_inference_steps=steps,
+        prompt=enhanced_prompt,
+        negative_prompt=enhanced_negative,
+        num_inference_steps=int(steps),
         guidance_scale=1.2,
-        width=resolution,
-        height=resolution,
+        width=int(resolution),
+        height=int(resolution),
         generator=gen
     ).images[0]
 
-    # 4️⃣ Progressive blur reveal
+    # 4️⃣ BLUR REVEAL
     max_blur = 20
-    steps_blur = 10
-    for i in range(steps_blur):
-        blur_percent = 100 - i*10
-        blurred_img = img.filter(ImageFilter.GaussianBlur(radius=max_blur * blur_percent / 100))
-        yield [blurred_img], "🟢 Generating..."
-
+    for i in range(10):
+        blur_pct = 100 - i*10
+        blurred = img.filter(ImageFilter.GaussianBlur(radius=max_blur * blur_pct/100))
+        yield [blurred], "🟢 Revealing..."
         time.sleep(1)
 
-    # 5️⃣ Final image with total time
-    total_time = round(time.time() - start_time, 2)
-    yield [img], f"✅ Done in {total_time}s"
+    # 5️⃣ FINAL
+    yield [img], f"✅ Done | Seed: {seed}"
 
 # -------------------------------
 # GRADIO UI
 # -------------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 👾 CREEPER AI — CPU SMART IMAGE GENERATION")
+    gr.Markdown("# 👾 CREEPER AI — SMART IMAGE GENERATOR")
+    gr.Markdown("1: The higher the resolution & steps, the longer the image takes to make.\n2: The more detailed the prompt and negative prompt, the better the result.")
 
     with gr.Row():
         with gr.Column():
-            prompt_in = gr.Textbox(label="Prompt", placeholder="cute snake", lines=2)
-
-            resolution = gr.Radio([256, 512, 768, 1024], value=512, label="Resolution")
+            prompt_in = gr.Textbox(label="Prompt")
+            neg_in = gr.Textbox(label="Negative Prompt")
 
-            steps = gr.Slider(6, 8, value=6, step=1, label="Steps")  # Minimum 6
+            resolution = gr.Radio([512, 768, 1024], value=512, label="Resolution")
+            steps = gr.Slider(6,8,value=6,step=1,label="Steps")
 
-            eta = gr.Markdown("⏱️ Estimated time: ~1m 0s")
+            eta = gr.Markdown("⏱️ Estimated: ~1m 0s")
             gen_btn = gr.Button("Generate")
 
             status = gr.Markdown("🟢 Ready")
@@ -116,15 +123,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column():
             gallery = gr.Gallery(columns=1)
 
-    # Update ETA dynamically
     for ctrl in [steps, resolution]:
         ctrl.change(estimate_time, [steps, resolution], eta)
 
-    # Button click triggers generation
     gen_btn.click(
         generate,
-        inputs=[prompt_in, resolution, steps],
+        inputs=[prompt_in, neg_in, resolution, steps],
         outputs=[gallery, status]
     )
 
-demo.launch()
+demo.launch()