Update app.py

app.py

@@ -1,124 +1,245 @@
 import gradio as gr
-import torch
-import random
-import time
 import requests
+import traceback
 
-from diffusers import DiffusionPipeline, LCMScheduler
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# ============================================================
+# -----------------------------
 # CONFIG
-# ============================================================
-SD_MODEL_ID = "runwayml/stable-diffusion-v1-5"
-LCM_ADAPTER_ID = "latent-consistency/lcm-lora-sdv1-5"
-TEXT_MODEL_ID = "HuggingFaceTB/SmolLM-135M-Instruct"
+# -----------------------------
+
+TEXT_MODEL = "HuggingFaceTB/SmolLM-135M-Instruct"
 
 Z_IMAGE_TURBO_API = "https://mrfakename-z-image-turbo.hf.space/gradio_api/mcp/"
 
-# ============================================================
-# LOAD IMAGE MODEL (CPU)
-# ============================================================
-pipe = DiffusionPipeline.from_pretrained(
-    SD_MODEL_ID,
-    safety_checker=None,
-    torch_dtype=torch.float32,
-)
-pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights(LCM_ADAPTER_ID)
-pipe.to("cpu")
-
-# ============================================================
-# LOAD TEXT MODEL (CPU)
-# ============================================================
-tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
-text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL_ID)
-text_model.to("cpu")
-
-# ============================================================
+DEFAULT_IMAGE_MODEL = "runwayml/stable-diffusion-v1-5"
+
+# -----------------------------
 # PROMPT ENHANCER
-# ============================================================
-def enhance_prompt(user_prompt: str) -> str:
-    instruction = (
+# -----------------------------
+
+def enhance_prompt(user_prompt):
+    system_prompt = (
         "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.\n\n"
+        "that requires clear instructions. Analyse the prompt, and output as "
+        "much detail as possible about it.\n\n"
         f"Prompt to enhance: {user_prompt}"
     )
 
-    inputs = tokenizer(instruction, return_tensors="pt")
+    from transformers import pipeline
+
+    pipe = pipeline(
+        "text-generation",
+        model=TEXT_MODEL,
+        max_new_tokens=500,
+        temperature=0.6,
+        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):
+    from diffusers import StableDiffusionPipeline
+    import torch
+
+    width, height = resolution
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        DEFAULT_IMAGE_MODEL,
+        torch_dtype=torch.float32
+    )
+
+    pipe = pipe.to("cpu")
+
+    image = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        num_inference_steps=20
+    ).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
+        )
+
+        if response.status_code != 200:
+            raise RuntimeError(
+                f"HTTP {response.status_code}: {response.text}"
+            )
 
-    with torch.no_grad():
-        output = text_model.generate(
-            **inputs,
-            max_new_tokens=500,
-            temperature=0.6,
-            do_sample=True,
+        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
 
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    # Return ONLY the model's generated enhancement
-    return decoded.split("Prompt to enhance:")[-1].strip()
+# -----------------------------
+# MAIN PIPELINE
+# -----------------------------
 
-# ============================================================
-# GENERATION FUNCTION
-# ============================================================
-def generate_image(
-    prompt,
+def generate(
+    user_prompt,
     negative_prompt,
-    resolution,
-    steps,
+    resolution_choice,
     ultra_speed,
+    ultra_steps
 ):
-    start_time = time.time()
+    width_map = {
+        "512": (512, 512),
+        "768": (768, 768),
+        "1024": (1024, 1024)
+    }
 
-    # White placeholder (always shown first)
-    placeholder = torch.ones((512, 512, 3))
-    yield [placeholder.numpy()], "🧠 Preparing..."
+    width, height = width_map[resolution_choice]
 
-    # ========================================================
-    # ULTRA SPEED PATH (EXTERNAL API)
-    # ========================================================
+    # ULTRA SPEED PATH
     if ultra_speed:
-        try:
-            payload = {
-                "prompt": prompt,
-                "width": int(resolution),
-                "height": int(resolution),
-                "num_inference_steps": int(steps),
-                "seed": random.randint(0, 10**9),
-                "randomize_seed": True,
-            }
-
-            response = requests.post(
-                Z_IMAGE_TURBO_API + "Z_Image_Turbo_generate_image",
-                json=payload,
-                timeout=120,
-            )
+        result = generate_ultra_image(
+            prompt=user_prompt,
+            width=width,
+            height=height,
+            steps=ultra_steps
+        )
+
+        if isinstance(result, tuple):
+            _, error = result
+            return None, error
 
-            if response.status_code != 200:
-                raise RuntimeError(
-                    f"HTTP {response.status_code}\n\nResponse body:\n{response.text}"
-                )
+        return result, ""
 
-            data = response.json()
+    # NORMAL PATH (TEXT → IMAGE)
+    enhanced_prompt = enhance_prompt(user_prompt)
 
-            if "images" not in data or len(data["images"]) == 0:
-                raise RuntimeError(
-                    f"Unexpected API response:\n{data}"
-                )
+    image = generate_cpu_image(
+        enhanced_prompt,
+        negative_prompt,
+        (width, height)
+    )
+
+    return image, ""
+
+
+# -----------------------------
+# UI
+# -----------------------------
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🖼️ Image Generator")
+
+    prompt = gr.Textbox(label="Prompt")
+    negative_prompt = gr.Textbox(label="Negative Prompt", value="")
+
+    ultra_speed = gr.Checkbox(
+        label="Ultra Speed (a few generations per day)",
+        value=False
+    )
+
+    resolution = gr.Radio(
+        ["512", "768", "1024"],
+        value="512",
+        label="Resolution"
+    )
 
-            image = data["images"][0]
+    ultra_steps = gr.Slider(
+        minimum=6,
+        maximum=20,
+        step=1,
+        value=8,
+        label="Ultra Speed Steps",
+        visible=False
+    )
+
+    generate_btn = gr.Button("Generate")
+
+    output_image = gr.Image()
+    error_box = gr.Markdown()
+
+    # -------------------------
+    # UI LOGIC
+    # -------------------------
+
+    def toggle_ultra(checked):
+        if checked:
+            return (
+                gr.Radio(
+                    ["512", "1024"],
+                    value="512"
+                ),
+                gr.update(visible=True)
+            )
+        else:
+            return (
+                gr.Radio(
+                    ["512", "768", "1024"],
+                    value="512"
+                ),
+                gr.update(visible=False)
+            )
 
-            elapsed = round(time.time() - start_time, 1)
-            yield [image], f"✅ Finished in {elapsed}s"
-            return
+    ultra_speed.change(
+        toggle_ultra,
+        inputs=ultra_speed,
+        outputs=[resolution, ultra_steps]
+    )
+
+    generate_btn.click(
+        generate,
+        inputs=[
+            prompt,
+            negative_prompt,
+            resolution,
+            ultra_speed,
+            ultra_steps
+        ],
+        outputs=[output_image, error_box]
+    )
 
-        except Exception as e:
-            error_message = f"""
-❌ **Ultra Speed Error**
 
-**Exception type:** `{type(e).__name__}`
+# -----------------------------
+# ENTRYPOINT
+# -----------------------------
 
-**Details:**
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)