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 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 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) decoded = tokenizer.decode(outputs[0], skip_special_tokens=True) 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.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) # =============================== # ULTRA MODE (Z-IMAGE-TURBO) # =============================== def call_ultra_api(prompt, negative, steps, resolution, seed): try: client = Client("mrfakename/Z-Image-Turbo") # The API doesn't support a separate 'negative_prompt' argument, # so we append it to the main prompt for better adherence. full_prompt = f"{prompt} . {negative}" if negative else prompt # Call the specific endpoint defined in your HTML result = client.predict( prompt=full_prompt, width=int(resolution), height=int(resolution), num_inference_steps=int(steps), seed=int(seed), randomize_seed=False, # We control the seed from Python api_name="/generate_image" ) # The API returns [image_filepath, seed_used] if isinstance(result, (list, tuple)): return Image.open(result[0]) return Image.open(result) except Exception as e: print(f"Ultra Mode Error: {e}") return None # =============================== # UI LOGIC # =============================== def toggle_ultra(is_ultra): if is_ultra: # Hide 768, extend steps to 20, show negative prompt return { negative_field: gr.update(visible=True), resolution: gr.update(choices=[512, 1024], value=512), steps: gr.update(minimum=1, maximum=20, value=9) } else: # Normal mode settings 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)) # Use user negative if provided, otherwise default final_neg = user_neg if user_neg.strip() else default_neg image = call_ultra_api(prompt, final_neg, step_val, size, seed) if image: yield (image, f"✅ Done (Ultra Mode). Seed: {seed}", gr.update(interactive=True)) else: yield (None, "❌ API Busy or Error. Try again.", gr.update(interactive=True)) else: # Normal CPU Mode 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() 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 * { color: black !important; } """ 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 - v4.0") gr.Markdown("Generate images using Creeper AI. Use **Ultra-Fast** for GPU speeds.") with gr.Row(): with gr.Column(scale=1): 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", 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") 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") # Connect the UI logic ultra_check.change( toggle_ultra, inputs=[ultra_check], outputs=[negative_field, resolution, steps] ) generate_btn.click( generate, inputs=[prompt_field, negative_field, resolution, steps, ultra_check], outputs=[output_img, status, generate_btn] ) demo.launch()