Upload generate_hf.py

generate_hf.py

@@ -0,0 +1,1193 @@
+"""
+Aniimage Generator — Generate anime images from text prompts.
+https://huggingface.co/8BitStudio/Aniimage-1
+
+Usage:
+    pip install torch torchvision diffusers transformers safetensors pillow huggingface_hub
+    python generate_hf.py
+"""
+
+import os
+import sys
+import torch
+import torch.nn.functional as F
+import numpy as np
+import tkinter as tk
+from tkinter import ttk, simpledialog
+from pathlib import Path
+from PIL import Image, ImageTk, ImageEnhance, ImageFilter
+from threading import Thread
+
+# ── Paths ─────────────────────────────────────────────────────────────────────
+SCRIPT_DIR = Path(__file__).resolve().parent
+MODEL_DIR = SCRIPT_DIR / "models"
+OUTPUT_DIR = SCRIPT_DIR / "generated"
+
+# ── HuggingFace repo ─────────────────────────────────────────────────────────
+HF_REPO_ID = "8BitStudio/Aniimage-1"
+
+# ── UNet config (must match training) ─────────────────────────────────────────
+UNET_CONFIG = dict(
+    sample_size=32,
+    in_channels=4,
+    out_channels=4,
+    block_out_channels=(256, 512, 768, 1024),
+    layers_per_block=2,
+    cross_attention_dim=768,
+    attention_head_dim=8,
+    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D",
+                      "CrossAttnDownBlock2D", "DownBlock2D"),
+    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D",
+                    "CrossAttnUpBlock2D", "UpBlock2D"),
+)
+
+VAE_ID = "stabilityai/sd-vae-ft-mse"
+CLIP_ID = "openai/clip-vit-large-patch14"
+
+SCHEDULER_LIST = [
+    "DPM++ 2M Karras",
+    "DPM++ SDE Karras",
+    "Euler a",
+    "Euler",
+    "DDIM",
+]
+
+DEFAULT_NEGATIVE = (
+    "low quality, ugly, blurry, distorted, deformed, bad anatomy, "
+    "bad proportions, extra limbs, missing limbs, watermark, text, "
+    "signature, washed out, flat colors, manga panel, disfigured, "
+    "poorly drawn, jpeg artifacts, cropped, out of frame"
+)
+
+
+# ── Model discovery ───────────────────────────────────────────────────────────
+
+def download_from_hf():
+    """Download model weights from HuggingFace if not already cached."""
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        print("Install huggingface_hub: pip install huggingface_hub")
+        return None
+
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    aniimage_dir = MODEL_DIR / "Aniimage-1"
+    weights_path = aniimage_dir / "diffusion_pytorch_model.safetensors"
+
+    if weights_path.exists():
+        print("Aniimage-1 weights already downloaded.")
+        return aniimage_dir
+
+    print(f"Downloading Aniimage-1 from {HF_REPO_ID}...")
+    aniimage_dir.mkdir(parents=True, exist_ok=True)
+
+    import shutil
+    dl_weights = hf_hub_download(repo_id=HF_REPO_ID,
+                                  filename="diffusion_pytorch_model.safetensors")
+    shutil.copy2(dl_weights, weights_path)
+
+    try:
+        dl_config = hf_hub_download(repo_id=HF_REPO_ID, filename="config.json")
+        shutil.copy2(dl_config, aniimage_dir / "config.json")
+    except Exception:
+        pass
+
+    print("Download complete!")
+    return aniimage_dir
+
+
+def find_models():
+    """Find all available models."""
+    options = []
+    if MODEL_DIR.exists():
+        for d in sorted(MODEL_DIR.iterdir()):
+            if d.is_dir():
+                safetensors = d / "diffusion_pytorch_model.safetensors"
+                ema_path = d / "ema_unet.pt"
+                unet_path = d / "unet.pt"
+                if safetensors.exists():
+                    options.append(("safetensors", d.name, d, "256"))
+                elif ema_path.exists() or unet_path.exists():
+                    options.append(("checkpoint", d.name, d, "256"))
+    return options
+
+
+# ── Theme ─────────────────────────────────────────────────────────────────────
+
+C = {
+    "bg":         "#111119",
+    "panel":      "#1b1b2f",
+    "card":       "#24243e",
+    "card_sel":   "#3a3a6e",
+    "border":     "#2e2e52",
+    "accent":     "#6c5ce7",
+    "accent_h":   "#8577ed",
+    "red":        "#e74c3c",
+    "green":      "#2ecc71",
+    "text":       "#eaeaea",
+    "text2":      "#a0a0b8",
+    "text3":      "#60607a",
+    "input":      "#16162a",
+    "input_fg":   "#dcdcf0",
+}
+
+
+class Generator:
+    def __init__(self, device="cuda"):
+        self.device = device if device == "cuda" and torch.cuda.is_available() else "cpu"
+        self.vae = None
+        self.text_encoder = None
+        self.tokenizer = None
+        self.unet = None
+        self.scheduler = None
+        self.loaded_checkpoint = None
+        self.latent_size = 32
+        self.output_size = 256
+        self.cancelled = False
+
+    def switch_device(self, new_device):
+        """Move all loaded models to a new device."""
+        new_device = new_device if new_device == "cuda" and torch.cuda.is_available() else "cpu"
+        if new_device == self.device:
+            return
+        self.device = new_device
+        if self.vae is not None:
+            self.vae = self.vae.to(self.device)
+        if self.text_encoder is not None:
+            self.text_encoder = self.text_encoder.to(self.device)
+        if self.unet is not None:
+            self.unet = self.unet.to(self.device)
+        self.loaded_checkpoint = None  # force reload on next generate
+        print(f"Switched to {self.device.upper()}")
+
+    def load_shared(self):
+        if self.vae is not None:
+            return
+        from diffusers import AutoencoderKL
+        from transformers import CLIPTextModel, CLIPTokenizer
+
+        print("Loading VAE...")
+        self.vae = AutoencoderKL.from_pretrained(VAE_ID).to(self.device)
+        self.vae.eval()
+
+        print("Loading CLIP text encoder...")
+        self.tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID)
+        self.text_encoder = CLIPTextModel.from_pretrained(CLIP_ID).to(self.device)
+        self.text_encoder.eval()
+
+        self.scheduler = self._make_scheduler("DPM++ 2M Karras")
+        self.scheduler_name = "DPM++ 2M Karras"
+        print("Shared models loaded.")
+
+    def _make_scheduler(self, name="DPM++ 2M Karras"):
+        from diffusers import (DDIMScheduler, DPMSolverMultistepScheduler,
+                               EulerAncestralDiscreteScheduler,
+                               EulerDiscreteScheduler)
+        base = dict(num_train_timesteps=1000, beta_schedule="scaled_linear",
+                    prediction_type="epsilon")
+        if name == "DPM++ 2M Karras":
+            return DPMSolverMultistepScheduler(
+                **base, algorithm_type="dpmsolver++",
+                solver_order=2, use_karras_sigmas=True)
+        elif name == "DPM++ SDE Karras":
+            return DPMSolverMultistepScheduler(
+                **base, algorithm_type="sde-dpmsolver++",
+                use_karras_sigmas=True)
+        elif name == "Euler a":
+            return EulerAncestralDiscreteScheduler(**base)
+        elif name == "Euler":
+            return EulerDiscreteScheduler(**base)
+        else:
+            return DDIMScheduler(**base, clip_sample=False,
+                                set_alpha_to_one=False)
+
+    def set_scheduler(self, name):
+        self.scheduler = self._make_scheduler(name)
+        self.scheduler_name = name
+
+    def load_model(self, model_path: Path, res_label: str = "256"):
+        if str(model_path) == self.loaded_checkpoint:
+            return
+        from diffusers import UNet2DConditionModel
+
+        self.load_shared()
+
+        if res_label == "512":
+            self.latent_size = 64
+            self.output_size = 512
+        else:
+            self.latent_size = 32
+            self.output_size = 256
+
+        unet_cfg = dict(UNET_CONFIG)
+        unet_cfg["sample_size"] = self.latent_size
+
+        print(f"Loading UNet from {model_path.name} ({res_label}px)...")
+        self.unet = UNet2DConditionModel(**unet_cfg).to(self.device)
+
+        safetensors_path = model_path / "diffusion_pytorch_model.safetensors"
+        ema_path = model_path / "ema_unet.pt"
+        unet_path = model_path / "unet.pt"
+
+        if safetensors_path.exists():
+            from safetensors.torch import load_file
+            state = load_file(str(safetensors_path), device=str(self.device))
+            self.unet.load_state_dict(state)
+            print("Loaded safetensors weights.")
+        elif ema_path.exists():
+            state = torch.load(ema_path, map_location=self.device, weights_only=True)
+            if "shadow_params" in state:
+                params = dict(self.unet.named_parameters())
+                keys = list(params.keys())
+                for i, sp in enumerate(state["shadow_params"]):
+                    params[keys[i]].data.copy_(sp)
+            else:
+                self.unet.load_state_dict(state)
+            print("Loaded EMA weights.")
+        elif unet_path.exists():
+            self.unet.load_state_dict(
+                torch.load(unet_path, map_location=self.device, weights_only=True))
+            print("Loaded UNet weights.")
+        else:
+            raise FileNotFoundError(f"No weights found in {model_path}")
+
+        self.unet.eval()
+        self.loaded_checkpoint = str(model_path)
+        print(f"Ready to generate at {self.output_size}x{self.output_size}!")
+
+    def _decode_latents(self, latents, post_process=False):
+        scaled = latents / self.vae.config.scaling_factor
+        with torch.no_grad():
+            image = self.vae.decode(scaled.float()).sample
+        image = (image.float() / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+        image = (image * 255).round().astype("uint8")
+        img = Image.fromarray(image)
+        if post_process:
+            img = self._post_process(img)
+        return img
+
+    def _sharpen_latents(self, latents, amount=0.08):
+        blurred = F.avg_pool2d(latents, kernel_size=3, stride=1, padding=1)
+        return latents + amount * (latents - blurred)
+
+    def _post_process(self, img):
+        img = img.filter(ImageFilter.UnsharpMask(radius=1.5, percent=40, threshold=2))
+        img = ImageEnhance.Contrast(img).enhance(1.06)
+        img = ImageEnhance.Color(img).enhance(1.10)
+        return img
+
+    def _image_quality_score(self, img: Image.Image) -> float:
+        arr = np.array(img.convert("L"), dtype=np.float32)
+        lap = (np.roll(arr, 1, 0) + np.roll(arr, -1, 0)
+             + np.roll(arr, 1, 1) + np.roll(arr, -1, 1) - 4.0 * arr)
+        sharpness = float(np.var(lap))
+        arr_rgb = np.array(img, dtype=np.float32)
+        color_var = float(np.mean(np.var(arr_rgb, axis=(0, 1))))
+        score = (sharpness * 0.6 + color_var * 0.4)
+        return min(100.0, score / 10.0)
+
+    @torch.no_grad()
+    def generate(self, prompt: str, negative_prompt: str = "",
+                 steps: int = 25, guidance_scale: float = 7.5,
+                 seed: int = -1, preview_callback=None,
+                 preview_every: int = 5) -> tuple:
+
+        if seed < 0:
+            seed = torch.randint(0, 2**32, (1,)).item()
+        gen = torch.Generator(device=self.device).manual_seed(seed)
+
+        tok = self.tokenizer(prompt, padding="max_length",
+                             max_length=self.tokenizer.model_max_length,
+                             truncation=True, return_tensors="pt")
+        text_emb = self.text_encoder(tok.input_ids.to(self.device))[0]
+
+        tok_neg = self.tokenizer(negative_prompt if negative_prompt else "",
+                                 padding="max_length",
+                                 max_length=self.tokenizer.model_max_length,
+                                 truncation=True, return_tensors="pt")
+        neg_emb = self.text_encoder(tok_neg.input_ids.to(self.device))[0]
+
+        text_emb_combined = torch.cat([neg_emb, text_emb])
+
+        scheduler = self._make_scheduler(self.scheduler_name)
+        scheduler.set_timesteps(steps, device=self.device)
+
+        latents = torch.randn(1, 4, self.latent_size, self.latent_size,
+                               generator=gen, device=self.device)
+        latents = latents * scheduler.init_noise_sigma
+
+        timesteps = scheduler.timesteps
+        total_steps = len(timesteps)
+
+        for step_i, t in enumerate(timesteps):
+            if self.cancelled:
+                return None, seed
+
+            latent_input = torch.cat([latents] * 2)
+            latent_input = scheduler.scale_model_input(latent_input, t)
+
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
+                                enabled=(self.device == "cuda")):
+                pred = self.unet(latent_input, t,
+                                 encoder_hidden_states=text_emb_combined).sample
+
+            pred_neg, pred_text = pred.chunk(2)
+            pred = pred_neg + guidance_scale * (pred_text - pred_neg)
+
+            latents = scheduler.step(pred, t, latents).prev_sample
+
+            if (preview_callback and step_i > 0
+                    and step_i % preview_every == 0
+                    and step_i < total_steps - 1):
+                preview = self._decode_latents(latents, post_process=False)
+                preview_callback(preview, step_i + 1, total_steps)
+
+        latents = self._sharpen_latents(latents)
+        final = self._decode_latents(latents, post_process=True)
+        return final, seed
+
+    @torch.no_grad()
+    def generate_adaptive(self, prompt: str, negative_prompt: str = "",
+                          base_steps: int = 25, max_steps: int = 85,
+                          guidance_scale: float = 7.5,
+                          quality_threshold: float = 45.0,
+                          preview_callback=None, preview_every: int = 5,
+                          status_callback=None) -> tuple:
+
+        result = self.generate(
+            prompt=prompt, negative_prompt=negative_prompt,
+            steps=base_steps, guidance_scale=guidance_scale,
+            preview_callback=preview_callback, preview_every=preview_every)
+
+        if result[0] is None:
+            return result
+
+        image, seed = result
+        quality = self._image_quality_score(image)
+
+        if status_callback:
+            status_callback(f"Quality: {quality:.1f}/100")
+
+        if quality >= quality_threshold:
+            return image, seed
+
+        rounds = 0
+        max_rounds = (max_steps - base_steps) // 20
+
+        while quality < quality_threshold and rounds < max_rounds:
+            if self.cancelled:
+                return image, seed
+            rounds += 1
+            if status_callback:
+                status_callback(f"Refining +20 steps (round {rounds})...")
+
+            refined = self.refine(
+                source_image=image, prompt=prompt,
+                negative_prompt=negative_prompt,
+                extra_steps=20, strength=0.3,
+                guidance_scale=guidance_scale,
+                preview_callback=preview_callback, preview_every=5)
+
+            if refined is None:
+                return image, seed
+            image = refined
+            quality = self._image_quality_score(image)
+
+            if status_callback:
+                status_callback(f"Quality after round {rounds}: {quality:.1f}/100")
+
+        return image, seed
+
+    @torch.no_grad()
+    def refine(self, source_image: Image.Image, prompt: str,
+               negative_prompt: str = "", extra_steps: int = 20,
+               strength: float = 0.35, guidance_scale: float = 7.5,
+               preview_callback=None, preview_every: int = 5) -> Image.Image:
+
+        img = source_image.resize((self.output_size, self.output_size), Image.LANCZOS)
+        img_tensor = torch.from_numpy(np.array(img)).float().div(127.5).sub(1.0)
+        img_tensor = img_tensor.permute(2, 0, 1).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            latents = self.vae.encode(img_tensor.float()).latent_dist.sample()
+        latents = latents * self.vae.config.scaling_factor
+
+        tok = self.tokenizer(prompt, padding="max_length",
+                             max_length=self.tokenizer.model_max_length,
+                             truncation=True, return_tensors="pt")
+        text_emb = self.text_encoder(tok.input_ids.to(self.device))[0]
+
+        tok_neg = self.tokenizer(negative_prompt if negative_prompt else "",
+                                 padding="max_length",
+                                 max_length=self.tokenizer.model_max_length,
+                                 truncation=True, return_tensors="pt")
+        neg_emb = self.text_encoder(tok_neg.input_ids.to(self.device))[0]
+        text_emb_combined = torch.cat([neg_emb, text_emb])
+
+        scheduler = self._make_scheduler(self.scheduler_name)
+        scheduler.set_timesteps(extra_steps, device=self.device)
+        start_step = max(0, int(len(scheduler.timesteps) * (1 - strength)))
+        timesteps = scheduler.timesteps[start_step:]
+
+        noise = torch.randn_like(latents)
+        latents = scheduler.add_noise(latents, noise, timesteps[:1])
+
+        total_steps = len(timesteps)
+        for step_i, t in enumerate(timesteps):
+            if self.cancelled:
+                return None
+            latent_input = torch.cat([latents] * 2)
+            latent_input = scheduler.scale_model_input(latent_input, t)
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
+                                enabled=(self.device == "cuda")):
+                pred = self.unet(latent_input, t,
+                                 encoder_hidden_states=text_emb_combined).sample
+            pred_neg, pred_text = pred.chunk(2)
+            pred = pred_neg + guidance_scale * (pred_text - pred_neg)
+            latents = scheduler.step(pred, t, latents).prev_sample
+
+            if (preview_callback and step_i > 0
+                    and step_i % preview_every == 0
+                    and step_i < total_steps - 1):
+                preview = self._decode_latents(latents, post_process=False)
+                preview_callback(preview, step_i + 1, total_steps)
+
+        latents = self._sharpen_latents(latents)
+        return self._decode_latents(latents, post_process=True)
+
+
+# ── GUI ───────────────────────────────────────────────────────────────────────
+
+class App:
+    def __init__(self):
+        self.gen = Generator()
+        self.models = find_models()
+        self.generated_images = []
+        self.generated_seeds = []
+        self.photo_refs = []
+        self.generating = False
+        self.selected_index = None
+
+        self.root = tk.Tk()
+        self.root.title("Aniimage")
+        self.root.configure(bg=C["bg"])
+        self.root.resizable(True, True)
+        self.root.geometry("900x780")
+        self.root.minsize(640, 500)
+
+        self._setup_styles()
+        self._build_ui()
+
+    def _setup_styles(self):
+        s = ttk.Style()
+        s.theme_use("clam")
+
+        # Base
+        s.configure(".", background=C["bg"], foreground=C["text"], font=("Segoe UI", 10))
+        s.configure("TFrame", background=C["bg"])
+        s.configure("TLabel", background=C["bg"], foreground=C["text"])
+        s.configure("TCheckbutton", background=C["bg"], foreground=C["text"])
+
+        # Combobox — readable text
+        s.configure("TCombobox", fieldbackground=C["input"], foreground=C["input_fg"],
+                    selectbackground=C["accent"], selectforeground="#ffffff",
+                    arrowcolor=C["text2"], padding=4)
+        s.map("TCombobox",
+              fieldbackground=[("readonly", C["input"])],
+              foreground=[("readonly", C["input_fg"])],
+              selectbackground=[("readonly", C["accent"])],
+              selectforeground=[("readonly", "#ffffff")])
+        # Combobox dropdown list colors
+        self.root.option_add("*TCombobox*Listbox.background", C["input"])
+        self.root.option_add("*TCombobox*Listbox.foreground", C["input_fg"])
+        self.root.option_add("*TCombobox*Listbox.selectBackground", C["accent"])
+        self.root.option_add("*TCombobox*Listbox.selectForeground", "#ffffff")
+        self.root.option_add("*TCombobox*Listbox.font", ("Segoe UI", 10))
+
+        # Spinbox
+        s.configure("TSpinbox", fieldbackground=C["input"], foreground=C["input_fg"],
+                    arrowcolor=C["text2"], padding=3)
+
+        # Buttons
+        s.configure("TButton", font=("Segoe UI", 10), padding=(14, 7),
+                    background=C["card"], foreground=C["text"])
+        s.map("TButton", background=[("active", C["card_sel"]), ("disabled", C["bg"])],
+              foreground=[("disabled", C["text3"])])
+
+        s.configure("Go.TButton", font=("Segoe UI", 11, "bold"), padding=(20, 9),
+                    background=C["accent"], foreground="#ffffff")
+        s.map("Go.TButton", background=[("active", C["accent_h"]),
+                                         ("disabled", C["border"])])
+
+        s.configure("Stop.TButton", font=("Segoe UI", 10, "bold"), padding=(14, 7),
+                    background=C["red"], foreground="#ffffff")
+        s.map("Stop.TButton", background=[("active", "#c0392b"),
+                                           ("disabled", C["border"])])
+
+        # Labelframe
+        s.configure("TLabelframe", background=C["bg"], foreground=C["text2"])
+        s.configure("TLabelframe.Label", background=C["bg"],
+                    foreground=C["text2"], font=("Segoe UI", 9, "bold"))
+
+        # Scrollbar
+        s.configure("Vertical.TScrollbar", background=C["card"],
+                    troughcolor=C["bg"], arrowcolor=C["text3"])
+
+    def _make_entry(self, parent, font_size=11, dim=False):
+        """Create a styled tk.Entry with readable text."""
+        return tk.Entry(parent, font=("Segoe UI", font_size),
+                        bg=C["input"], fg=C["input_fg"] if not dim else C["text2"],
+                        insertbackground=C["input_fg"],
+                        relief="flat", bd=6,
+                        selectbackground=C["accent"], selectforeground="#ffffff",
+                        highlightthickness=1, highlightcolor=C["accent"],
+                        highlightbackground=C["border"])
+
+    def _build_ui(self):
+        # ── Header ────────────────────────────────────────────────────────
+        header = tk.Frame(self.root, bg=C["panel"], padx=20, pady=12)
+        header.pack(fill=tk.X)
+
+        tk.Label(header, text="Aniimage", bg=C["panel"], fg=C["accent"],
+                 font=("Segoe UI", 20, "bold")).pack(side=tk.LEFT)
+        tk.Label(header, text="by 8BitStudio", bg=C["panel"], fg=C["text3"],
+                 font=("Segoe UI", 10)).pack(side=tk.LEFT, padx=(10, 0), pady=(6, 0))
+
+        # Device switch — right side of header
+        device_frame = tk.Frame(header, bg=C["panel"])
+        device_frame.pack(side=tk.RIGHT)
+
+        tk.Label(device_frame, text="Device:", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9)).pack(side=tk.LEFT, padx=(0, 5))
+
+        self.device_var = tk.StringVar(value="GPU" if self.gen.device == "cuda" else "CPU")
+        devices = ["GPU", "CPU"] if torch.cuda.is_available() else ["CPU"]
+        device_combo = ttk.Combobox(device_frame, textvariable=self.device_var,
+                                     values=devices, state="readonly", width=5)
+        device_combo.pack(side=tk.LEFT)
+        device_combo.bind("<<ComboboxSelected>>", self._on_device_change)
+
+        # ── Main content — two-column: controls left, images right ────────
+        main = tk.Frame(self.root, bg=C["bg"])
+        main.pack(fill=tk.BOTH, expand=True, padx=12, pady=(8, 12))
+
+        # Left panel (controls)
+        left = tk.Frame(main, bg=C["panel"], width=340, padx=16, pady=12)
+        left.pack(side=tk.LEFT, fill=tk.Y, padx=(0, 8))
+        left.pack_propagate(False)
+
+        # Right panel (image grid)
+        right = tk.Frame(main, bg=C["bg"])
+        right.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
+
+        self._build_controls(left)
+        self._build_grid(right)
+
+    def _build_controls(self, parent):
+        # ── Model ─────────────────────────────────────────────────────────
+        tk.Label(parent, text="Model", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9, "bold")).pack(anchor=tk.W)
+
+        self.model_var = tk.StringVar()
+        model_names = [m[1] for m in self.models] or ["No models found"]
+        self.model_combo = ttk.Combobox(parent, textvariable=self.model_var,
+                                        values=model_names, state="readonly", width=32)
+        self.model_combo.pack(fill=tk.X, pady=(3, 12))
+        self.model_combo.current(len(model_names) - 1)
+
+        # ── Prompt ────────────────────────────────────────────────────────
+        tk.Label(parent, text="Prompt", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9, "bold")).pack(anchor=tk.W)
+        self.prompt_entry = self._make_entry(parent)
+        self.prompt_entry.pack(fill=tk.X, pady=(3, 8))
+        self.prompt_entry.insert(0, "a smiling anime girl with long blue hair")
+        self.prompt_entry.bind("<Return>", lambda e: self.on_generate())
+
+        # ── Negative prompt ───────────────────────────────────────────────
+        tk.Label(parent, text="Negative prompt", bg=C["panel"], fg=C["text3"],
+                 font=("Segoe UI", 9)).pack(anchor=tk.W)
+        self.neg_entry = self._make_entry(parent, font_size=9, dim=True)
+        self.neg_entry.pack(fill=tk.X, pady=(3, 12))
+        self.neg_entry.insert(0, DEFAULT_NEGATIVE)
+
+        # ── Settings grid ─────────────────────────────────────────────────
+        grid = tk.Frame(parent, bg=C["panel"])
+        grid.pack(fill=tk.X, pady=(0, 8))
+
+        # Row 1: Scheduler
+        tk.Label(grid, text="Scheduler", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9)).grid(row=0, column=0, sticky="w", pady=(0, 6))
+        self.scheduler_var = tk.StringVar(value="DPM++ 2M Karras")
+        sched_combo = ttk.Combobox(grid, textvariable=self.scheduler_var,
+                                    values=SCHEDULER_LIST, state="readonly", width=18)
+        sched_combo.grid(row=0, column=1, columnspan=3, sticky="ew", padx=(8, 0), pady=(0, 6))
+        sched_combo.bind("<<ComboboxSelected>>", self._on_scheduler_change)
+
+        # Row 2: Steps, CFG, Count
+        tk.Label(grid, text="Steps", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9)).grid(row=1, column=0, sticky="w", pady=(0, 6))
+        self.steps_var = tk.StringVar(value="25")
+        tk.Entry(grid, textvariable=self.steps_var, width=5, font=("Segoe UI", 10),
+                 bg=C["input"], fg=C["input_fg"], insertbackground=C["input_fg"],
+                 relief="flat", bd=4).grid(row=1, column=1, sticky="w", padx=(8, 12), pady=(0, 6))
+
+        tk.Label(grid, text="CFG", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9)).grid(row=1, column=2, sticky="w", pady=(0, 6))
+        self.cfg_var = tk.StringVar(value="7.5")
+        tk.Entry(grid, textvariable=self.cfg_var, width=5, font=("Segoe UI", 10),
+                 bg=C["input"], fg=C["input_fg"], insertbackground=C["input_fg"],
+                 relief="flat", bd=4).grid(row=1, column=3, sticky="w", padx=(8, 0), pady=(0, 6))
+
+        # Row 3: Count, Live preview
+        tk.Label(grid, text="Count", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9)).grid(row=2, column=0, sticky="w", pady=(0, 6))
+        self.count_var = tk.StringVar(value="4")
+        ttk.Spinbox(grid, from_=1, to=12, textvariable=self.count_var, width=4,
+                     font=("Segoe UI", 10)).grid(row=2, column=1, sticky="w", padx=(8, 12), pady=(0, 6))
+
+        self.live_preview_var = tk.BooleanVar(value=False)
+        ttk.Checkbutton(grid, text="Live preview",
+                         variable=self.live_preview_var).grid(
+            row=2, column=2, columnspan=2, sticky="w", pady=(0, 6))
+
+        grid.columnconfigure(1, weight=1)
+        grid.columnconfigure(3, weight=1)
+
+        # ── Auto quality ──────────────────────────────────────────────────
+        self.auto_quality_var = tk.BooleanVar(value=False)
+        ttk.Checkbutton(parent, text="Auto quality (refine if undercooked)",
+                         variable=self.auto_quality_var).pack(anchor=tk.W, pady=(0, 12))
+
+        # ── Buttons ───────────────────────────────────────────────────────
+        btn_frame = tk.Frame(parent, bg=C["panel"])
+        btn_frame.pack(fill=tk.X, pady=(0, 10))
+
+        self.gen_btn = ttk.Button(btn_frame, text="Generate", command=self.on_generate,
+                                  style="Go.TButton")
+        self.gen_btn.pack(fill=tk.X, pady=(0, 5))
+
+        btn_row = tk.Frame(btn_frame, bg=C["panel"])
+        btn_row.pack(fill=tk.X)
+
+        self.stop_btn = ttk.Button(btn_row, text="Stop", command=self.on_stop,
+                                   state=tk.DISABLED, style="Stop.TButton")
+        self.stop_btn.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 3))
+
+        self.save_btn = ttk.Button(btn_row, text="Save Selected", command=self.on_save,
+                                   state=tk.DISABLED)
+        self.save_btn.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(3, 3))
+
+        self.save_all_btn = ttk.Button(btn_row, text="Save All", command=self.on_save_all,
+                                       state=tk.DISABLED)
+        self.save_all_btn.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(3, 0))
+
+        # ── Prompt queue ─────────────────────────────────────────────────
+        sep = tk.Frame(parent, height=1, bg=C["border"])
+        sep.pack(fill=tk.X, pady=(8, 10))
+
+        tk.Label(parent, text="Prompt Queue", bg=C["panel"], fg=C["text2"],
+                 font=("Segoe UI", 9, "bold")).pack(anchor=tk.W)
+
+        queue_input = tk.Frame(parent, bg=C["panel"])
+        queue_input.pack(fill=tk.X, pady=(4, 0))
+
+        self.queue_entry = self._make_entry(queue_input, font_size=9)
+        self.queue_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 4))
+        self.queue_entry.bind("<Return>", lambda e: self._queue_add())
+
+        ttk.Button(queue_input, text="Add", width=4,
+                   command=self._queue_add).pack(side=tk.LEFT)
+
+        self.queue_listbox = tk.Listbox(
+            parent, height=4, bg=C["input"], fg=C["input_fg"],
+            selectbackground=C["accent"], selectforeground="#fff",
+            font=("Segoe UI", 9), activestyle="none",
+            relief="flat", bd=4, highlightthickness=0)
+        self.queue_listbox.pack(fill=tk.X, pady=(5, 0))
+
+        queue_btns = tk.Frame(parent, bg=C["panel"])
+        queue_btns.pack(fill=tk.X, pady=(4, 0))
+
+        self.queue_run_btn = ttk.Button(queue_btns, text="Run Queue",
+                                         command=self.on_run_queue, style="Go.TButton")
+        self.queue_run_btn.pack(side=tk.LEFT, padx=(0, 4))
+
+        for txt, cmd in [("Remove", self._queue_remove), ("Clear", self._queue_clear),
+                         ("Up", self._queue_move_up), ("Down", self._queue_move_down),
+                         ("+ Current", self._queue_add_current)]:
+            ttk.Button(queue_btns, text=txt, command=cmd).pack(side=tk.LEFT, padx=2)
+
+        # ── Status bar ────────────────────────────────────────────────────
+        status_frame = tk.Frame(parent, bg=C["bg"], padx=8, pady=6)
+        status_frame.pack(fill=tk.X, side=tk.BOTTOM)
+
+        self.status_var = tk.StringVar(value="Ready")
+        tk.Label(status_frame, textvariable=self.status_var,
+                 bg=C["bg"], fg=C["green"], font=("Segoe UI", 9),
+                 anchor="w").pack(fill=tk.X)
+
+    def _build_grid(self, parent):
+        self.canvas = tk.Canvas(parent, bg=C["bg"], highlightthickness=0)
+        scrollbar = ttk.Scrollbar(parent, orient=tk.VERTICAL, command=self.canvas.yview)
+        self.grid_frame = tk.Frame(self.canvas, bg=C["bg"])
+
+        self.grid_frame.bind("<Configure>",
+                              lambda e: self.canvas.configure(
+                                  scrollregion=self.canvas.bbox("all")))
+        self.canvas_window = self.canvas.create_window((0, 0), window=self.grid_frame,
+                                                        anchor="nw")
+        self.canvas.configure(yscrollcommand=scrollbar.set)
+
+        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
+        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
+
+        self.canvas.bind("<Configure>", self._on_canvas_resize)
+        self.canvas.bind_all("<MouseWheel>",
+                              lambda e: self.canvas.yview_scroll(
+                                  int(-1 * (e.delta / 120)), "units"))
+
+        self.placeholder = tk.Label(self.grid_frame,
+                                     text="Generated images\nwill appear here",
+                                     bg=C["bg"], fg=C["text3"],
+                                     font=("Segoe UI", 13), justify="center")
+        self.placeholder.grid(row=0, column=0, pady=80)
+
+    # ── Event handlers ────────────────────────────────────────────────────
+
+    def _on_device_change(self, event=None):
+        choice = self.device_var.get()
+        new_dev = "cuda" if choice == "GPU" else "cpu"
+        self.status_var.set(f"Switching to {choice}...")
+        self.root.update()
+        self.gen.switch_device(new_dev)
+        self.status_var.set(f"Now using {choice}")
+
+    def _on_scheduler_change(self, event=None):
+        name = self.scheduler_var.get()
+        self.gen.set_scheduler(name)
+        self.status_var.set(f"Scheduler: {name}")
+
+    def _on_canvas_resize(self, event):
+        self.canvas.itemconfig(self.canvas_window, width=event.width)
+        if self.generated_images:
+            self._layout_grid()
+
+    def _get_grid_cols(self):
+        canvas_w = self.canvas.winfo_width()
+        if canvas_w < 50:
+            canvas_w = 560
+        tile_size = self._get_tile_size()
+        return max(1, canvas_w // (tile_size + 16))
+
+    def _get_tile_size(self):
+        n = len(self.generated_images)
+        if n <= 2: return 260
+        elif n <= 4: return 220
+        elif n <= 6: return 180
+        else: return 160
+
+    def _layout_grid(self):
+        for w in self.grid_frame.winfo_children():
+            w.destroy()
+        self.photo_refs.clear()
+
+        if not self.generated_images:
+            return
+
+        tile_size = self._get_tile_size()
+        cols = self._get_grid_cols()
+
+        for i, (img, seed) in enumerate(zip(self.generated_images, self.generated_seeds)):
+            row, col = divmod(i, cols)
+            is_selected = (i == self.selected_index)
+
+            card_bg = C["accent"] if is_selected else C["card"]
+            card = tk.Frame(self.grid_frame, bg=card_bg, padx=3, pady=3)
+            card.grid(row=row, column=col, padx=5, pady=5, sticky="nsew")
+
+            display = img.resize((tile_size, tile_size), Image.LANCZOS)
+            photo = ImageTk.PhotoImage(display)
+            self.photo_refs.append(photo)
+
+            img_label = tk.Label(card, image=photo, bg=card_bg, bd=0)
+            img_label.pack()
+            img_label.bind("<Button-1>", lambda e, idx=i: self._select_image(idx))
+            img_label.bind("<Button-3>", lambda e, idx=i: self._show_refine_menu(e, idx))
+
+            tk.Label(card, text=f"seed: {seed}", bg=card_bg,
+                     fg=C["text3"], font=("Segoe UI", 8)).pack()
+
+        for c in range(cols):
+            self.grid_frame.columnconfigure(c, weight=1)
+
+    def _select_image(self, idx):
+        if idx >= len(self.generated_images):
+            return
+        self.selected_index = idx
+        self.save_btn.configure(state=tk.NORMAL)
+        self.status_var.set(f"Selected image {idx + 1} (seed: {self.generated_seeds[idx]})")
+        self._layout_grid()
+
+    def _show_refine_menu(self, event, idx):
+        if self.generating:
+            return
+        menu = tk.Menu(self.root, tearoff=0, bg=C["card"], fg=C["text"],
+                       activebackground=C["accent"], activeforeground="#fff",
+                       font=("Segoe UI", 10), bd=0)
+        menu.add_command(label="  Refine (more steps)...  ",
+                         command=lambda: self._ask_refine(idx))
+        menu.tk_popup(event.x_root, event.y_root)
+
+    def _ask_refine(self, idx):
+        extra = simpledialog.askinteger(
+            "Refine Image", "Extra denoising steps:",
+            initialvalue=20, minvalue=5, maxvalue=200, parent=self.root)
+        if extra is None:
+            return
+        self._select_image(idx)
+        self.generating = True
+        self.gen.cancelled = False
+        self.gen_btn.configure(state=tk.DISABLED)
+        self.stop_btn.configure(state=tk.NORMAL)
+        self.status_var.set(f"Refining image {idx + 1}...")
+        self.root.update()
+        Thread(target=self._refine_thread, args=(idx, extra), daemon=True).start()
+
+    def _refine_thread(self, idx, extra_steps):
+        try:
+            source = self.generated_images[idx]
+            prompt = self.prompt_entry.get().strip()
+            neg = self.neg_entry.get().strip()
+            cfg = float(self.cfg_var.get())
+            callback = self._show_preview if self.live_preview_var.get() else None
+
+            refined = self.gen.refine(
+                source_image=source, prompt=prompt, negative_prompt=neg,
+                extra_steps=extra_steps, guidance_scale=cfg,
+                preview_callback=callback, preview_every=5)
+
+            if refined is not None:
+                self.generated_images[idx] = refined
+                self.generated_seeds[idx] = f"{self.generated_seeds[idx]}+R{extra_steps}"
+                self._layout_grid()
+                self.status_var.set(f"Refined image {idx + 1}")
+            else:
+                self.status_var.set("Refine stopped.")
+            self.root.update()
+        except Exception as e:
+            self.status_var.set(f"Refine error: {e}")
+            import traceback; traceback.print_exc()
+        finally:
+            self.generating = False
+            self.gen.cancelled = False
+            self.gen_btn.configure(state=tk.NORMAL)
+            self.stop_btn.configure(state=tk.DISABLED)
+
+    # ── Queue ─────────────────────────────────────────────────────────────
+
+    def _queue_add(self):
+        text = self.queue_entry.get().strip()
+        if text:
+            self.queue_listbox.insert(tk.END, text)
+            self.queue_entry.delete(0, tk.END)
+
+    def _queue_add_current(self):
+        text = self.prompt_entry.get().strip()
+        if text:
+            self.queue_listbox.insert(tk.END, text)
+
+    def _queue_remove(self):
+        sel = self.queue_listbox.curselection()
+        if sel:
+            self.queue_listbox.delete(sel[0])
+
+    def _queue_clear(self):
+        self.queue_listbox.delete(0, tk.END)
+
+    def _queue_move_up(self):
+        sel = self.queue_listbox.curselection()
+        if sel and sel[0] > 0:
+            idx = sel[0]
+            text = self.queue_listbox.get(idx)
+            self.queue_listbox.delete(idx)
+            self.queue_listbox.insert(idx - 1, text)
+            self.queue_listbox.selection_set(idx - 1)
+
+    def _queue_move_down(self):
+        sel = self.queue_listbox.curselection()
+        if sel and sel[0] < self.queue_listbox.size() - 1:
+            idx = sel[0]
+            text = self.queue_listbox.get(idx)
+            self.queue_listbox.delete(idx)
+            self.queue_listbox.insert(idx + 1, text)
+            self.queue_listbox.selection_set(idx + 1)
+
+    def on_run_queue(self):
+        if self.generating or not self.models:
+            return
+        prompts = list(self.queue_listbox.get(0, tk.END))
+        if not prompts:
+            self.status_var.set("Queue is empty")
+            return
+        self.generating = True
+        self.gen.cancelled = False
+        self.gen_btn.configure(state=tk.DISABLED)
+        self.queue_run_btn.configure(state=tk.DISABLED)
+        self.stop_btn.configure(state=tk.NORMAL)
+        Thread(target=self._queue_thread, args=(prompts,), daemon=True).start()
+
+    def _queue_thread(self, prompts):
+        try:
+            idx = self.model_combo.current()
+            mdl = self.models[idx]
+            self.status_var.set(f"Loading {mdl[1]}...")
+            self.root.update()
+            self.gen.load_model(mdl[2], mdl[3])
+
+            neg = self.neg_entry.get().strip()
+            steps = int(self.steps_var.get())
+            cfg = float(self.cfg_var.get())
+            num_images = max(1, min(12, int(self.count_var.get())))
+            live_preview = self.live_preview_var.get()
+            auto_quality = self.auto_quality_var.get()
+
+            self.generated_images.clear()
+            self.generated_seeds.clear()
+            self.selected_index = None
+            if self.placeholder:
+                self.placeholder.destroy()
+                self.placeholder = None
+
+            for p_idx, prompt in enumerate(prompts):
+                if self.gen.cancelled:
+                    break
+                self.queue_listbox.selection_clear(0, tk.END)
+                self.queue_listbox.selection_set(p_idx)
+                self.queue_listbox.see(p_idx)
+
+                for img_i in range(num_images):
+                    if self.gen.cancelled:
+                        break
+                    self.status_var.set(
+                        f"[{p_idx + 1}/{len(prompts)}] image {img_i + 1}/{num_images}")
+                    self.root.update()
+
+                    callback = None
+                    if live_preview:
+                        self._setup_preview_card()
+                        callback = self._show_preview
+
+                    if auto_quality:
+                        image, used_seed = self.gen.generate_adaptive(
+                            prompt=prompt, negative_prompt=neg,
+                            base_steps=steps, max_steps=steps + 60,
+                            guidance_scale=cfg,
+                            preview_callback=callback, preview_every=5,
+                            status_callback=lambda m: (
+                                self.status_var.set(m), self.root.update()))
+                    else:
+                        image, used_seed = self.gen.generate(
+                            prompt=prompt, negative_prompt=neg,
+                            steps=steps, guidance_scale=cfg,
+                            preview_callback=callback, preview_every=5)
+
+                    if image is None:
+                        break
+                    self.generated_images.append(image)
+                    self.generated_seeds.append(used_seed)
+                    save_path = self._next_save_path(prompt)
+                    image.save(save_path)
+                    self._layout_grid()
+                    self.root.update()
+
+                if self.gen.cancelled:
+                    break
+
+            done = len(self.generated_images)
+            self.status_var.set(
+                f"Queue {'stopped' if self.gen.cancelled else 'done'}! {done} images saved.")
+            if done > 0:
+                self.save_all_btn.configure(state=tk.NORMAL)
+
+        except Exception as e:
+            self.status_var.set(f"Queue error: {e}")
+            import traceback; traceback.print_exc()
+        finally:
+            self.generating = False
+            self.gen.cancelled = False
+            self.gen_btn.configure(state=tk.NORMAL)
+            self.queue_run_btn.configure(state=tk.NORMAL)
+            self.stop_btn.configure(state=tk.DISABLED)
+
+    # ── Generation ────────────────────────────────────────────────────────
+
+    def on_stop(self):
+        if self.generating:
+            self.gen.cancelled = True
+            self.status_var.set("Stopping...")
+            self.root.update()
+
+    def on_generate(self):
+        if self.generating or not self.models:
+            return
+        self.generating = True
+        self.gen.cancelled = False
+        self.gen_btn.configure(state=tk.DISABLED)
+        self.stop_btn.configure(state=tk.NORMAL)
+        self.status_var.set("Loading model...")
+        self.root.update()
+        Thread(target=self._generate_thread, daemon=True).start()
+
+    def _setup_preview_card(self):
+        tile_size = self._get_tile_size()
+        cols = self._get_grid_cols()
+        row, col = divmod(len(self.generated_images), cols)
+        card = tk.Frame(self.grid_frame, bg=C["card"], padx=3, pady=3)
+        card.grid(row=row, column=col, padx=5, pady=5, sticky="nsew")
+        self._preview_label = tk.Label(card, bg=C["card"],
+                                        width=tile_size, height=tile_size)
+        self._preview_label.pack()
+        self.root.update()
+
+    def _show_preview(self, preview_img, step, total):
+        tile_size = self._get_tile_size()
+        display = preview_img.resize((tile_size, tile_size), Image.LANCZOS)
+        photo = ImageTk.PhotoImage(display)
+        self._preview_photo = photo
+        if hasattr(self, '_preview_label') and self._preview_label.winfo_exists():
+            self._preview_label.configure(image=photo)
+        self.status_var.set(f"Step {step}/{total}")
+        self.root.update()
+
+    def _generate_thread(self):
+        try:
+            idx = self.model_combo.current()
+            mdl = self.models[idx]
+            self.status_var.set(f"Loading {mdl[1]}...")
+            self.root.update()
+            self.gen.load_model(mdl[2], mdl[3])
+
+            prompt = self.prompt_entry.get().strip()
+            neg = self.neg_entry.get().strip()
+            steps = int(self.steps_var.get())
+            cfg = float(self.cfg_var.get())
+            num_images = max(1, min(12, int(self.count_var.get())))
+            live_preview = self.live_preview_var.get()
+            auto_quality = self.auto_quality_var.get()
+
+            self.generated_images.clear()
+            self.generated_seeds.clear()
+            self.selected_index = None
+            if self.placeholder:
+                self.placeholder.destroy()
+                self.placeholder = None
+
+            for i in range(num_images):
+                if self.gen.cancelled:
+                    break
+                self.status_var.set(f"Generating {i + 1}/{num_images}...")
+                self.root.update()
+
+                callback = None
+                if live_preview:
+                    self._setup_preview_card()
+                    callback = self._show_preview
+
+                if auto_quality:
+                    image, used_seed = self.gen.generate_adaptive(
+                        prompt=prompt, negative_prompt=neg,
+                        base_steps=steps, max_steps=steps + 60,
+                        guidance_scale=cfg,
+                        preview_callback=callback, preview_every=5,
+                        status_callback=lambda m: (
+                            self.status_var.set(m), self.root.update()))
+                else:
+                    image, used_seed = self.gen.generate(
+                        prompt=prompt, negative_prompt=neg,
+                        steps=steps, guidance_scale=cfg,
+                        preview_callback=callback, preview_every=5)
+
+                if image is None:
+                    break
+                self.generated_images.append(image)
+                self.generated_seeds.append(used_seed)
+                self._layout_grid()
+                self.root.update()
+
+            done = len(self.generated_images)
+            if self.gen.cancelled:
+                self.status_var.set(f"Stopped. {done} image(s) kept.")
+            else:
+                self.status_var.set(f"Done! {done} images. Click to select.")
+            if done > 0:
+                self.save_all_btn.configure(state=tk.NORMAL)
+            self.save_btn.configure(state=tk.DISABLED)
+
+        except Exception as e:
+            self.status_var.set(f"Error: {e}")
+            import traceback; traceback.print_exc()
+        finally:
+            self.generating = False
+            self.gen.cancelled = False
+            self.gen_btn.configure(state=tk.NORMAL)
+            self.stop_btn.configure(state=tk.DISABLED)
+
+    # ── Save ──────────────────────────────────────────────────────────────
+
+    def _next_save_path(self, prompt_text):
+        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+        slug = prompt_text.strip()[:50] if prompt_text.strip() else "untitled"
+        base = OUTPUT_DIR / f"{slug}.png"
+        if not base.exists():
+            return base
+        n = 1
+        while True:
+            path = OUTPUT_DIR / f"{slug} {n}.png"
+            if not path.exists():
+                return path
+            n += 1
+
+    def on_save(self):
+        if self.selected_index is None or not self.generated_images:
+            return
+        img = self.generated_images[self.selected_index]
+        path = self._next_save_path(self.prompt_entry.get().strip())
+        img.save(path)
+        self.status_var.set(f"Saved: {path.name}")
+
+    def on_save_all(self):
+        if not self.generated_images:
+            return
+        prompt_text = self.prompt_entry.get().strip()
+        for img in self.generated_images:
+            path = self._next_save_path(prompt_text)
+            img.save(path)
+        self.status_var.set(f"Saved {len(self.generated_images)} images")
+
+    def run(self):
+        self.root.mainloop()
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    models = find_models()
+    if not models:
+        print("No models found locally. Downloading from HuggingFace...")
+        result = download_from_hf()
+        if result:
+            models = find_models()
+
+    if not models:
+        print("No models found!")
+        print(f"Place model weights in: {MODEL_DIR}/YourModelName/")
+        print("Expected files: diffusion_pytorch_model.safetensors or ema_unet.pt")
+        sys.exit(1)
+
+    print(f"Found {len(models)} model(s): {', '.join(m[1] for m in models)}")
+    print(f"Device: {'CUDA (GPU)' if torch.cuda.is_available() else 'CPU'}")
+    print("Starting Aniimage...")
+
+    app = App()
+    app.run()