Update app.py

app.py

@@ -2,171 +2,161 @@ import gradio as gr
 import torch
 import random
 import time
+import io
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from diffusers import DiffusionPipeline, LCMScheduler
 from PIL import Image, ImageFilter
+from gradio_client import Client
 
+# ===============================
+# LOCAL MODELS (CPU MODE)
+# ===============================
 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:
     if not user_prompt.strip():
         return "A beautiful digital painting of a fantasy landscape"
-        
     instruction = (
-        f"<|im_start|>system\nYou are a prompt engineer. Expand the user's prompt into a detailed, descriptive visual prompt for SD 1.5. Output only the enhanced prompt.<|im_end|>\n"
+        f"<|im_start|>system\nYou are a prompt engineer. Expand the user's prompt into a detailed visual prompt. Output only the enhanced prompt.<|im_end|>\n"
         f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
     inputs = tokenizer(instruction, return_tensors="pt")
     with torch.no_grad():
-        outputs = text_model.generate(
-            **inputs,
-            max_new_tokens=500,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
+        outputs = text_model.generate(**inputs, max_new_tokens=500, temperature=0.7, do_sample=True, pad_token_id=tokenizer.eos_token_id)
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if "assistant" in decoded:
-        decoded = decoded.split("assistant")[-1]
-
-    return decoded.strip()
+    return decoded.split("assistant")[-1].strip() if "assistant" in decoded else decoded.strip()
 
 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 = 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)
 
-def estimate_time(steps, res):
-    res = int(res)
-    steps = int(steps)
-    base_sec = {512: 15, 768: 35, 1024: 75}[res]
-    total = (steps * base_sec) + 10
-    return f"⏱️ **Estimate**: ~{total//60}m {total%60}s"
-
-def generate(prompt, negative, resolution, steps):
-    size = int(resolution)
-    
-    yield (
-        None, 
-        "🧠 Analysing Prompt",
-        gr.update(interactive=False)
-    )
-
-    enhanced = enhance_prompt(prompt)
-    
-    yield (
-        None,
-        "🎨 Generating Image...",
-        gr.update(interactive=False)
-    )
-
-    seed = random.randint(0, 2**32 - 1)
-    generator = torch.Generator("cpu").manual_seed(seed)
-
-    start = time.time()
-    image = pipe(
-        prompt=enhanced,
-        negative_prompt=negative,
-        num_inference_steps=int(steps),
-        guidance_scale=1.0,
-        width=size,
-        height=size,
-        generator=generator
-    ).images[0]
-    elapsed = int(time.time() - start)
-
-    for i in range(5):
-        blur_radius = (5 - i) * 2
-        blur = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
-        yield (
-            blur,
-            "🎨 Generating Image...",
-            gr.update(interactive=False)
+# ===============================
+# ULTRA MODE (REMOTE API)
+# ===============================
+def call_ultra_api(prompt, negative, steps, seed):
+    try:
+        client = Client("mrfakename/Z-Image-Turbo")
+        result = client.predict(
+            prompt=prompt,
+            negative_prompt=negative,
+            num_inference_steps=steps,
+            seed=seed,
+            api_name="/predict"
         )
-        time.sleep(0.2)
-
-    yield (
-        image,
-        f"✅ Done in {elapsed}s.",
-        gr.update(interactive=True)
-    )
-
+        # The client returns a path to the image
+        return Image.open(result)
+    except Exception as e:
+        print(f"API Error: {e}")
+        return None
+
+# ===============================
+# UI LOGIC
+# ===============================
+def toggle_ultra(is_ultra):
+    if is_ultra:
+        return {
+            negative_field: gr.update(visible=True),
+            resolution: gr.update(choices=[512, 1024], value=512),
+            steps: gr.update(minimum=6, maximum=20, value=6)
+        }
+    else:
+        return {
+            negative_field: gr.update(visible=False),
+            resolution: gr.update(choices=[512, 768, 1024], value=512),
+            steps: gr.update(minimum=6, maximum=10, value=6)
+        }
+
+def generate(prompt, user_neg, res, step_val, is_ultra):
+    size = int(res)
+    seed = random.randint(0, 2**32 - 1)
+    default_neg = "blurry, low quality, distorted, watermark"
+    
+    if is_ultra:
+        yield (None, "🎨 Generating Image...", gr.update(interactive=False))
+        # Ultra mode uses user_neg (field is visible) and skips enhancement
+        image = call_ultra_api(prompt, user_neg, step_val, seed)
+        if image:
+            yield (image, f"✅ Done in Ultra Mode.", gr.update(interactive=True))
+        else:
+            yield (None, "❌ API Error or Limit Reached", gr.update(interactive=True))
+    else:
+        yield (None, "🧠 Analysing Prompt", gr.update(interactive=False))
+        enhanced = enhance_prompt(prompt)
+        
+        yield (None, "🎨 Generating Image...", gr.update(interactive=False))
+        generator = torch.Generator("cpu").manual_seed(seed)
+        start = time.time()
+        
+        # Normal mode uses default_neg
+        image = pipe(
+            prompt=enhanced, 
+            negative_prompt=default_neg, 
+            num_inference_steps=int(step_val), 
+            guidance_scale=1.0, 
+            width=size, 
+            height=size, 
+            generator=generator
+        ).images[0]
+        
+        elapsed = int(time.time() - start)
+        for i in range(5):
+            blur = image.filter(ImageFilter.GaussianBlur(radius=(5 - i) * 2))
+            yield (blur, "🎨 Generating Image...", gr.update(interactive=False))
+            time.sleep(0.2)
+        yield (image, f"✅ Done in {elapsed}s.", gr.update(interactive=True))
+
+# ===============================
+# INTERFACE
+# ===============================
 custom_css = """
 #container { max-width: 1000px; margin: auto; }
 .generate-btn { background: linear-gradient(90deg, #2ecc71, #27ae60) !important; color: white !important; }
-.status-box { 
-    font-size: 1.1em; 
-    padding: 10px; 
-    border-radius: 8px; 
-    background: #ffffff !important; 
-    border: 1px solid #ddd;
-}
+.status-box { font-size: 1.1em; padding: 10px; border-radius: 8px; background: #ffffff !important; border: 1px solid #ddd; }
 .status-box * { color: black !important; }
 """
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="green", secondary_hue="gray"), css=custom_css) as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="green"), css=custom_css) as demo:
     with gr.Column(elem_id="container"):
-        gr.Markdown("# 👾 Creeper AI - v3.3")
+        gr.Markdown("# 👾 Creeper AI - v4.0")
         gr.Markdown("Generate images using Creeper AI, running on CPU using fast models.")
         
         with gr.Row():
             with gr.Column(scale=1):
-                prompt = gr.Textbox(
-                    label="What do you want to see?", 
-                    placeholder="e.g. A futuristic city in a forest",
-                    lines=3
-                )
-                negative = gr.Textbox(
-                    label="Negative Prompt", 
-                    value="blurry, low quality, distorted, watermark",
-                    placeholder="Things to avoid..."
-                )
+                prompt_field = gr.Textbox(label="What do you want to see?", placeholder="e.g. A futuristic city", lines=3)
+                negative_field = gr.Textbox(label="Negative Prompt", value="blurry, low quality, distorted, watermark", visible=False)
                 
                 with gr.Accordion("Settings ⚙️", open=True):
-                    resolution = gr.Radio(
-                        [512, 768, 1024], 
-                        value=512, 
-                        label="Resolution"
-                    )
-                    steps = gr.Slider(
-                        6, 10, 
-                        value=6, 
-                        step=1, 
-                        label="Inference Steps"
-                    )
-                
-                eta = gr.Markdown("⏱️ **Estimate**: ~1m 40s")
+                    resolution = gr.Radio([512, 768, 1024], value=512, label="Resolution")
+                    steps = gr.Slider(6, 10, value=6, step=1, label="Inference Steps")
+                    ultra_check = gr.Checkbox(label="Ultra-Fast Gen (A few IMGs per day)")
+
                 generate_btn = gr.Button("🚀 Generate Image", variant="primary", elem_classes="generate-btn")
 
             with gr.Column(scale=1):
                 output_img = gr.Image(label="Result", interactive=False)
                 status = gr.Markdown("🟢 Ready", elem_classes="status-box")
 
-    resolution.change(estimate_time, [steps, resolution], eta)
-    steps.change(estimate_time, [steps, resolution], eta)
+    # Connect the UI logic
+    ultra_check.change(
+        toggle_ultra, 
+        inputs=[ultra_check], 
+        outputs=[negative_field, resolution, steps]
+    )
 
     generate_btn.click(
         generate,
-        inputs=[prompt, negative, resolution, steps],
+        inputs=[prompt_field, negative_field, resolution, steps, ultra_check],
         outputs=[output_img, status, generate_btn]
     )
 
-demo.launch()
+demo.launch()