Patrick Plate
353 lines
15 KiB
Python
353 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
FLUX Image Generator — Desktop GUI
|
|
Supports: FLUX.1 Schnell (fast) and FLUX.2 Klein Heretic (unrestricted)
|
|
Run: python3 gui.py
|
|
No external dependencies — stdlib + tkinter only.
|
|
"""
|
|
|
|
import json
|
|
import random
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
from pathlib import Path
|
|
import tkinter as tk
|
|
from tkinter import ttk, scrolledtext, messagebox
|
|
|
|
COMFYUI = "http://localhost:8188"
|
|
OUTPUT_DIR = Path.home() / "Pictures" / "mcp-generated"
|
|
WORKFLOWS_DIR = Path(__file__).parent / "src/workflows"
|
|
|
|
# ── Model definitions ────────────────────────────────────────────────────────
|
|
# Each entry defines how to patch the workflow for that model.
|
|
# node_pos / node_neg: CLIPTextEncode node IDs for positive/negative prompts
|
|
# node_latent: latent image node (for width/height)
|
|
# node_seed: where to write the seed value
|
|
# node_save: SaveImage node (for filename_prefix)
|
|
# node_steps: dict of {node_id: field} for steps — None if not patchable
|
|
MODELS = {
|
|
"FLUX.2 Klein Heretic (unrestricted)": {
|
|
"workflow": "flux2_klein_heretic.json",
|
|
"default_steps": 20,
|
|
"node_pos": ("2", "text"),
|
|
"node_neg": ("3", "text"),
|
|
"node_latent": ("6", "width", "height"),
|
|
"node_seed": ("10", "noise_seed"),
|
|
"node_steps": ("7", "steps"),
|
|
"node_save": ("13", "filename_prefix"),
|
|
"description": "FLUX.2 Klein 4B + Heretic abliterated encoder. ~50s/image. No refusals.",
|
|
},
|
|
"FLUX.1 Schnell (fast)": {
|
|
"workflow": "flux_schnell.json",
|
|
"default_steps": 4,
|
|
"node_pos": ("6", "text"),
|
|
"node_neg": ("33", "text"),
|
|
"node_latent": ("27", "width", "height"),
|
|
"node_seed": ("13", "seed"), # KSampler has seed directly
|
|
"node_steps": ("13", "steps"),
|
|
"node_save": ("9", "filename_prefix"),
|
|
"description": "FLUX.1 Schnell — fast (~5s/image), standard quality. Has safety filter.",
|
|
},
|
|
}
|
|
|
|
PRESETS = {
|
|
"Square 1024": (1024, 1024),
|
|
"Landscape 16:9": (1280, 720),
|
|
"Portrait 9:16": (720, 1280),
|
|
"Wide 3:2": (1536, 1024),
|
|
"Tall 2:3": (1024, 1536),
|
|
"BFL Wide 7:4": (1344, 768),
|
|
"BFL Tall 4:7": (768, 1344),
|
|
}
|
|
|
|
|
|
def load_workflow(filename):
|
|
with open(WORKFLOWS_DIR / filename) as f:
|
|
return json.load(f)
|
|
|
|
|
|
def patch_workflow(wf, spec, prompt, neg, width, height, steps, seed, name):
|
|
"""Apply generation parameters to a workflow dict in-place."""
|
|
node_pos_id, node_pos_field = spec["node_pos"]
|
|
node_neg_id, node_neg_field = spec["node_neg"]
|
|
lat_id, lat_w, lat_h = spec["node_latent"]
|
|
seed_id, seed_field = spec["node_seed"]
|
|
save_id, save_field = spec["node_save"]
|
|
|
|
wf[node_pos_id]["inputs"][node_pos_field] = prompt
|
|
wf[node_neg_id]["inputs"][node_neg_field] = neg
|
|
wf[lat_id]["inputs"][lat_w] = width
|
|
wf[lat_id]["inputs"][lat_h] = height
|
|
wf[seed_id]["inputs"][seed_field] = seed
|
|
wf[save_id]["inputs"][save_field] = name
|
|
|
|
if spec["node_steps"]:
|
|
steps_id, steps_field = spec["node_steps"]
|
|
wf[steps_id]["inputs"][steps_field] = steps
|
|
|
|
return wf
|
|
|
|
|
|
def submit_prompt(workflow):
|
|
data = json.dumps({"prompt": workflow}).encode()
|
|
req = urllib.request.Request(
|
|
f"{COMFYUI}/prompt", data=data,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=30) as resp:
|
|
return json.loads(resp.read())["prompt_id"]
|
|
|
|
|
|
def wait_for_image(prompt_id, timeout=300):
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
with urllib.request.urlopen(f"{COMFYUI}/history/{prompt_id}", timeout=10) as resp:
|
|
history = json.loads(resp.read())
|
|
if prompt_id in history:
|
|
for node_out in history[prompt_id].get("outputs", {}).values():
|
|
if "images" in node_out:
|
|
return node_out["images"][0]
|
|
return None
|
|
time.sleep(2)
|
|
raise TimeoutError("Timed out waiting for image")
|
|
|
|
|
|
def download_image(filename, subfolder=""):
|
|
url = f"{COMFYUI}/view?filename={filename}&subfolder={subfolder}&type=output"
|
|
with urllib.request.urlopen(url, timeout=30) as resp:
|
|
return resp.read()
|
|
|
|
|
|
class App(tk.Tk):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.title("FLUX Image Generator")
|
|
self.resizable(True, True)
|
|
self.minsize(760, 640)
|
|
self._current_image = None
|
|
self._build_ui()
|
|
|
|
def _build_ui(self):
|
|
main = ttk.Frame(self, padding=12)
|
|
main.pack(fill="both", expand=True)
|
|
main.columnconfigure(0, weight=1)
|
|
main.columnconfigure(1, weight=2)
|
|
|
|
# ── LEFT PANEL ───────────────────────────────────────────────────────
|
|
left = ttk.Frame(main)
|
|
left.grid(row=0, column=0, sticky="nsew", padx=(0, 8))
|
|
left.columnconfigure(0, weight=1)
|
|
row = 0
|
|
|
|
# Model selector
|
|
ttk.Label(left, text="Model", font=("", 10, "bold")).grid(
|
|
row=row, column=0, sticky="w"); row += 1
|
|
self.model_var = tk.StringVar(value=list(MODELS.keys())[0])
|
|
model_cb = ttk.Combobox(left, textvariable=self.model_var,
|
|
values=list(MODELS.keys()), state="readonly", width=40)
|
|
model_cb.grid(row=row, column=0, sticky="ew", pady=(2, 2)); row += 1
|
|
model_cb.bind("<<ComboboxSelected>>", self._on_model_change)
|
|
self.model_desc = ttk.Label(left, text="", foreground="gray",
|
|
wraplength=300, font=("", 8))
|
|
self.model_desc.grid(row=row, column=0, sticky="w", pady=(0, 8)); row += 1
|
|
# description updated after all widgets are created (see end of _build_ui)
|
|
|
|
# Prompt
|
|
ttk.Label(left, text="Prompt", font=("", 10, "bold")).grid(
|
|
row=row, column=0, sticky="w"); row += 1
|
|
self.prompt_txt = scrolledtext.ScrolledText(left, height=6, wrap="word", font=("", 10))
|
|
self.prompt_txt.grid(row=row, column=0, sticky="ew", pady=(2, 8)); row += 1
|
|
|
|
# Negative prompt
|
|
ttk.Label(left, text="Negative Prompt (optional)").grid(
|
|
row=row, column=0, sticky="w"); row += 1
|
|
self.neg_txt = scrolledtext.ScrolledText(left, height=3, wrap="word", font=("", 9))
|
|
self.neg_txt.grid(row=row, column=0, sticky="ew", pady=(2, 8)); row += 1
|
|
|
|
# Size preset
|
|
ttk.Label(left, text="Size Preset").grid(row=row, column=0, sticky="w"); row += 1
|
|
self.preset_var = tk.StringVar(value="Square 1024")
|
|
preset_cb = ttk.Combobox(left, textvariable=self.preset_var,
|
|
values=list(PRESETS.keys()), state="readonly")
|
|
preset_cb.grid(row=row, column=0, sticky="ew", pady=(2, 2)); row += 1
|
|
preset_cb.bind("<<ComboboxSelected>>", self._apply_preset)
|
|
|
|
# Width / Height
|
|
wh = ttk.Frame(left)
|
|
wh.grid(row=row, column=0, sticky="ew", pady=(4, 8)); row += 1
|
|
wh.columnconfigure(1, weight=1)
|
|
wh.columnconfigure(3, weight=1)
|
|
ttk.Label(wh, text="W").grid(row=0, column=0, padx=(0, 4))
|
|
self.width_var = tk.IntVar(value=1024)
|
|
ttk.Spinbox(wh, from_=256, to=4096, increment=8, textvariable=self.width_var,
|
|
width=6).grid(row=0, column=1, sticky="ew")
|
|
ttk.Label(wh, text="H", padding=(8, 0, 4, 0)).grid(row=0, column=2)
|
|
self.height_var = tk.IntVar(value=1024)
|
|
ttk.Spinbox(wh, from_=256, to=4096, increment=8, textvariable=self.height_var,
|
|
width=6).grid(row=0, column=3, sticky="ew")
|
|
|
|
# Steps
|
|
steps_row = ttk.Frame(left)
|
|
steps_row.grid(row=row, column=0, sticky="ew", pady=(0, 4)); row += 1
|
|
steps_row.columnconfigure(1, weight=1)
|
|
ttk.Label(steps_row, text="Steps").grid(row=0, column=0, padx=(0, 8))
|
|
self.steps_var = tk.IntVar(value=20)
|
|
self.steps_lbl = ttk.Label(steps_row, text="20", width=3)
|
|
self.steps_lbl.grid(row=0, column=2, padx=(6, 0))
|
|
ttk.Scale(steps_row, from_=1, to=60, variable=self.steps_var,
|
|
orient="horizontal",
|
|
command=lambda v: self.steps_lbl.config(text=str(int(float(v))))
|
|
).grid(row=0, column=1, sticky="ew")
|
|
|
|
# Count
|
|
count_row = ttk.Frame(left)
|
|
count_row.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1
|
|
count_row.columnconfigure(1, weight=1)
|
|
ttk.Label(count_row, text="Count").grid(row=0, column=0, padx=(0, 8))
|
|
self.count_var = tk.IntVar(value=1)
|
|
ttk.Spinbox(count_row, from_=1, to=20, textvariable=self.count_var,
|
|
width=4).grid(row=0, column=1, sticky="w")
|
|
|
|
# Seed
|
|
seed_row = ttk.Frame(left)
|
|
seed_row.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1
|
|
seed_row.columnconfigure(1, weight=1)
|
|
ttk.Label(seed_row, text="Seed").grid(row=0, column=0, padx=(0, 8))
|
|
self.seed_var = tk.StringVar(value="-1")
|
|
ttk.Entry(seed_row, textvariable=self.seed_var, width=12).grid(row=0, column=1, sticky="w")
|
|
ttk.Button(seed_row, text="🎲", width=3,
|
|
command=lambda: self.seed_var.set(str(random.randint(0, 2**32 - 1)))
|
|
).grid(row=0, column=2, padx=(4, 0))
|
|
ttk.Label(seed_row, text="(-1 = random)", foreground="gray"
|
|
).grid(row=0, column=3, padx=(4, 0))
|
|
|
|
# Name
|
|
name_row = ttk.Frame(left)
|
|
name_row.grid(row=row, column=0, sticky="ew", pady=(0, 12)); row += 1
|
|
name_row.columnconfigure(1, weight=1)
|
|
ttk.Label(name_row, text="Name").grid(row=0, column=0, padx=(0, 8))
|
|
self.name_var = tk.StringVar(value="img")
|
|
ttk.Entry(name_row, textvariable=self.name_var).grid(row=0, column=1, sticky="ew")
|
|
|
|
# Generate button
|
|
self.gen_btn = ttk.Button(left, text="⚡ Generate", command=self._start_generation)
|
|
self.gen_btn.grid(row=row, column=0, sticky="ew", pady=(0, 8)); row += 1
|
|
|
|
# Status + progress
|
|
self.status_var = tk.StringVar(value="Ready")
|
|
ttk.Label(left, textvariable=self.status_var, foreground="gray",
|
|
wraplength=300).grid(row=row, column=0, sticky="w"); row += 1
|
|
self.progress = ttk.Progressbar(left, mode="indeterminate")
|
|
self.progress.grid(row=row, column=0, sticky="ew", pady=(4, 0)); row += 1
|
|
|
|
# ── RIGHT PANEL — preview ────────────────────────────────────────────
|
|
right = ttk.LabelFrame(main, text="Preview", padding=8)
|
|
right.grid(row=0, column=1, sticky="nsew")
|
|
right.columnconfigure(0, weight=1)
|
|
right.rowconfigure(0, weight=1)
|
|
|
|
self.preview_lbl = ttk.Label(right, text="No image yet", anchor="center",
|
|
background="#1a1a1a", foreground="#888")
|
|
self.preview_lbl.grid(row=0, column=0, sticky="nsew")
|
|
|
|
self.path_lbl = ttk.Label(right, text="", foreground="gray", font=("", 8))
|
|
self.path_lbl.grid(row=1, column=0, sticky="w", pady=(4, 0))
|
|
|
|
ttk.Button(right, text="Open folder",
|
|
command=self._open_folder).grid(row=2, column=0, sticky="e", pady=(4, 0))
|
|
|
|
# Init model description + steps default now that all widgets exist
|
|
self._on_model_change()
|
|
|
|
def _on_model_change(self, _=None):
|
|
spec = MODELS[self.model_var.get()]
|
|
self.model_desc.config(text=spec["description"])
|
|
self.steps_var.set(spec["default_steps"])
|
|
self.steps_lbl.config(text=str(spec["default_steps"]))
|
|
|
|
def _apply_preset(self, _=None):
|
|
w, h = PRESETS[self.preset_var.get()]
|
|
self.width_var.set(w)
|
|
self.height_var.set(h)
|
|
|
|
def _start_generation(self):
|
|
prompt = self.prompt_txt.get("1.0", "end").strip()
|
|
if not prompt:
|
|
messagebox.showwarning("No prompt", "Please enter a prompt.")
|
|
return
|
|
self.gen_btn.config(state="disabled")
|
|
self.progress.start(10)
|
|
self.status_var.set("Generating…")
|
|
t = threading.Thread(target=self._run_generation, args=(prompt,), daemon=True)
|
|
t.start()
|
|
|
|
def _run_generation(self, prompt):
|
|
try:
|
|
neg = self.neg_txt.get("1.0", "end").strip()
|
|
steps = int(self.steps_var.get())
|
|
width = int(self.width_var.get())
|
|
height = int(self.height_var.get())
|
|
count = int(self.count_var.get())
|
|
name = self.name_var.get().strip() or "img"
|
|
seed_str = self.seed_var.get().strip()
|
|
base_seed = int(seed_str) if seed_str else -1
|
|
|
|
model_name = self.model_var.get()
|
|
spec = MODELS[model_name]
|
|
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
for i in range(count):
|
|
seed = (base_seed if base_seed == -1 else base_seed + i)
|
|
if seed == -1:
|
|
seed = random.randint(0, 2**32 - 1)
|
|
|
|
label = f"{name}_{i+1:02d}" if count > 1 else name
|
|
self.after(0, self.status_var.set,
|
|
f"[{i+1}/{count}] {model_name.split('(')[0].strip()} · seed {seed}…")
|
|
|
|
wf = load_workflow(spec["workflow"])
|
|
patch_workflow(wf, spec, prompt, neg, width, height, steps, seed, label)
|
|
|
|
prompt_id = submit_prompt(wf)
|
|
img_info = wait_for_image(prompt_id)
|
|
|
|
if img_info:
|
|
img_data = download_image(img_info["filename"], img_info.get("subfolder", ""))
|
|
out_path = OUTPUT_DIR / f"{label}_{seed}.png"
|
|
out_path.write_bytes(img_data)
|
|
self.after(0, self._show_preview, out_path)
|
|
|
|
self.after(0, self.status_var.set,
|
|
f"✅ Done — {count} image(s) saved to ~/Pictures/mcp-generated/")
|
|
except Exception as exc:
|
|
self.after(0, self.status_var.set, f"❌ Error: {exc}")
|
|
finally:
|
|
self.after(0, self.progress.stop)
|
|
self.after(0, lambda: self.gen_btn.config(state="normal"))
|
|
|
|
def _show_preview(self, path):
|
|
try:
|
|
photo = tk.PhotoImage(file=str(path))
|
|
pw, ph = photo.width(), photo.height()
|
|
subsample = 1
|
|
while pw // subsample > 600 or ph // subsample > 600:
|
|
subsample += 1
|
|
if subsample > 1:
|
|
photo = photo.subsample(subsample, subsample)
|
|
self.preview_lbl.config(image=photo, text="")
|
|
self._current_image = photo
|
|
self.path_lbl.config(text=str(path))
|
|
except Exception as e:
|
|
self.status_var.set(f"Preview error: {e}")
|
|
|
|
def _open_folder(self):
|
|
import subprocess
|
|
subprocess.Popen(["xdg-open", str(OUTPUT_DIR)])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app = App()
|
|
app.mainloop()
|