Update app.py

app.py

@@ -1,228 +1,156 @@
 import gradio as gr
-import requests
-import traceback
-
-# -----------------------------
-# CONFIG
-# -----------------------------
-TEXT_MODEL = "HuggingFaceTB/SmolLM-135M-Instruct"
-Z_IMAGE_TURBO_API = "https://mrfakename-z-image-turbo.hf.space/gradio_api/mcp/"
-DEFAULT_IMAGE_MODEL = "runwayml/stable-diffusion-v1-5"
-
-# -----------------------------
-# PROMPT ENHANCER
-# -----------------------------
-def enhance_prompt(user_prompt):
-    system_prompt = (
+import torch
+import random
+import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from diffusers import DiffusionPipeline, LCMScheduler
+from PIL import Image, ImageFilter
+
+# ===============================
+# 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:
+    instruction = (
         "Please enhance this prompt so it is suitable for an image generator "
         "that requires clear instructions. Analyse the prompt, and output as "
-        "much detail as possible about it.\n\n"
-        f"Prompt to enhance: {user_prompt}"
+        "much visual detail as possible about it.\n\n"
+        f"Prompt to enhance: {user_prompt}\n\n"
+        "Enhanced prompt:"
     )
 
-    from transformers import pipeline
-
-    pipe = pipeline(
-        "text-generation",
-        model=TEXT_MODEL,
+    inputs = tokenizer(instruction, return_tensors="pt")
+    outputs = text_model.generate(
+        **inputs,
         max_new_tokens=500,
         temperature=0.6,
-        do_sample=True
+        do_sample=True,
     )
 
-    result = pipe(system_prompt)
-    return result[0]["generated_text"].strip()
-
-# -----------------------------
-# DEFAULT CPU IMAGE GENERATOR
-# -----------------------------
-def generate_cpu_image(prompt, negative_prompt, resolution, steps):
-    from diffusers import StableDiffusionPipeline
-    import torch
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    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,
+    torch_dtype=torch.float32,
+    safety_checker=None
+)
+
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.load_lora_weights(LCM_LORA)
+pipe.to("cpu")
+
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+pipe.set_progress_bar_config(disable=True)
+
+# ===============================
+# TIME ESTIMATION
+# ===============================
+def estimate_time(steps, res):
+    res = int(res)
+    steps = int(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(prompt, negative, resolution, steps):
+    size = int(resolution)
+
+    # Placeholder while thinking
+    yield (
+        [Image.new("RGB", (size, size), "white")],
+        "🧠 Enhancing prompt..."
+    )
 
-    width, height = resolution
+    enhanced = enhance_prompt(prompt)
 
-    pipe = StableDiffusionPipeline.from_pretrained(
-        DEFAULT_IMAGE_MODEL,
-        torch_dtype=torch.float32
+    yield (
+        [Image.new("RGB", (size, size), "white")],
+        "🎨 Generating image..."
     )
 
-    pipe = pipe.to("cpu")
+    seed = random.randint(0, 1_000_000_000)
+    generator = torch.Generator("cpu").manual_seed(seed)
 
+    pipe.scheduler.set_timesteps(int(steps))
+
+    start = time.time()
     image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height,
-        num_inference_steps=steps
+        prompt=enhanced,
+        negative_prompt=negative,
+        num_inference_steps=int(steps),
+        guidance_scale=1.2,
+        width=size,
+        height=size,
+        generator=generator
     ).images[0]
-
-    return image
-
-# -----------------------------
-# ULTRA SPEED IMAGE GENERATOR
-# -----------------------------
-def generate_ultra_image(prompt, width, height, steps):
-    payload = {
-        "prompt": prompt,
-        "width": width,
-        "height": height,
-        "num_inference_steps": steps,
-        "seed": 0,
-        "randomize_seed": True
-    }
-
-    try:
-        response = requests.post(
-            Z_IMAGE_TURBO_API + "Z_Image_Turbo_generate_image",
-            json=payload,
-            timeout=120
+    elapsed = int(time.time() - start)
+
+    # Blur reveal
+    for i in range(10):
+        blur = image.filter(ImageFilter.GaussianBlur(radius=(10 - i)))
+        yield (
+            [blur],
+            f"👀 Revealing... ({i+1}/10)"
         )
+        time.sleep(1)
 
-        if response.status_code != 200:
-            raise RuntimeError(
-                f"HTTP {response.status_code}: {response.text}"
-            )
-
-        data = response.json()
-
-        if "error" in data:
-            raise RuntimeError(str(data["error"]))
-
-        return data["image"]
-
-    except Exception as e:
-        error_message = (
-            "❌ **Ultra Speed Error**\n\n"
-            f"**Exception type:** `{type(e).__name__}`\n\n"
-            "**Details:**\n"
-            "```\n"
-            f"{str(e)}\n"
-            "```\n\n"
-            "**Traceback:**\n"
-            "```\n"
-            f"{traceback.format_exc()}\n"
-            "```"
-        )
-        return None, error_message
-
-# -----------------------------
-# MAIN PIPELINE
-# -----------------------------
-def generate(
-    user_prompt,
-    negative_prompt,
-    resolution_choice,
-    ultra_speed,
-    steps_slider
-):
-    width_map = {
-        "512": (512, 512),
-        "768": (768, 768),
-        "1024": (1024, 1024)
-    }
-
-    width, height = width_map[resolution_choice]
-
-    # ULTRA SPEED PATH
-    if ultra_speed:
-        result = generate_ultra_image(
-            prompt=user_prompt,
-            width=width,
-            height=height,
-            steps=steps_slider
-        )
-
-        if isinstance(result, tuple):
-            _, error = result
-            return None, error
-
-        return result, ""
-
-    # NORMAL PATH (TEXT → IMAGE)
-    enhanced_prompt = enhance_prompt(user_prompt)
-
-    image = generate_cpu_image(
-        enhanced_prompt,
-        negative_prompt,
-        (width, height),
-        steps_slider
+    yield (
+        [image],
+        f"✅ Done in {elapsed}s | Seed {seed}"
     )
 
-    return image, ""
-
-# -----------------------------
+# ===============================
 # UI
-# -----------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("## 👾 CREEPER AI IMAGE GENERATOR")
+# ===============================
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    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️⃣ The more detailed the prompt and negative prompt, the better the result."
     )
 
-    prompt = gr.Textbox(label="Prompt")
-    negative_prompt = gr.Textbox(label="Negative Prompt", value="")
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt")
+            negative = gr.Textbox(label="Negative Prompt")
 
-    ultra_speed = gr.Checkbox(
-        label="Ultra Speed (a few generations per day)",
-        value=False
-    )
+            resolution = gr.Radio([512, 768, 1024], value=512, label="Resolution")
+            steps = gr.Slider(6, 8, value=6, step=1, label="Steps")
 
-    resolution = gr.Radio(
-        ["512", "768", "1024"],
-        value="512",
-        label="Resolution"
-    )
+            eta = gr.Markdown("⏱️ Estimated: ~1m")
+            generate_btn = gr.Button("Generate")
 
-    steps_slider = gr.Slider(
-        minimum=6,
-        maximum=8,
-        step=1,
-        value=6,
-        label="Steps"
-    )
+            status = gr.Markdown("🟢 Ready")
 
-    generate_btn = gr.Button("Generate")
-    output_image = gr.Image()
-    error_box = gr.Markdown()
-
-    # -------------------------
-    # UI LOGIC
-    # -------------------------
-    def toggle_ultra(checked):
-        # Update resolution choices and step range instantly
-        if checked:
-            return (
-                gr.update(choices=["512", "1024"], value="512"),
-                gr.update(minimum=6, maximum=20, value=6)
-            )
-        else:
-            return (
-                gr.update(choices=["512", "768", "1024"], value="512"),
-                gr.update(minimum=6, maximum=8, value=6)
-            )
-
-    ultra_speed.change(
-        toggle_ultra,
-        inputs=ultra_speed,
-        outputs=[resolution, steps_slider]
-    )
+        with gr.Column():
+            gallery = gr.Gallery(columns=1)
+
+    resolution.change(estimate_time, [steps, resolution], eta)
+    steps.change(estimate_time, [steps, resolution], eta)
 
     generate_btn.click(
         generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            resolution,
-            ultra_speed,
-            steps_slider
-        ],
-        outputs=[output_image, error_box]
+        inputs=[prompt, negative, resolution, steps],
+        outputs=[gallery, status]
     )
 
-# -----------------------------
-# ENTRYPOINT
-# -----------------------------
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(share=False)