Upload generate_images.py

generate_images.py

@@ -0,0 +1,1459 @@
+"""
+Aniimage Generator — Generate anime images from text prompts.
+https://huggingface.co/8BitStudio/Aniimage-2
+
+Usage:
+    pip install -U torch torchvision "diffusers>=0.37.1" "transformers>=4.46,<5" accelerate safetensors pillow huggingface_hub
+    python generate_hf_aniimage2_corrected.py
+"""
+
+import os
+import sys
+import gc
+import json
+import torch
+import numpy as np
+import tkinter as tk
+from tkinter import ttk, simpledialog
+from pathlib import Path
+from PIL import Image, ImageTk
+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-2"
+
+# ── Aniimage-2 training configuration fallback ────────────────────────────────
+# The downloaded model_config.json is preferred. These values mirror it so the
+# launcher still behaves correctly if only the UNet files were copied locally.
+UNET_CONFIG = dict(
+    sample_size=64,
+    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=("CrossAttnDownBlock2D", "CrossAttnDownBlock2D",
+                      "CrossAttnDownBlock2D", "DownBlock2D"),
+    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D",
+                    "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+)
+
+# Aniimage-2 was trained with this VAE, not the SD 1.x MSE VAE.
+VAE_ID = "madebyollin/sdxl-vae-fp16-fix"
+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 _read_json(path: Path):
+    """Read a JSON file, returning an empty dict when it is unusable."""
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return data if isinstance(data, dict) else {}
+    except (OSError, ValueError, TypeError):
+        return {}
+
+
+def _looks_like_unet_config(config: dict) -> bool:
+    """Return True when a config contains the core Diffusers UNet fields."""
+    required = {
+        "in_channels", "out_channels", "block_out_channels",
+        "down_block_types", "up_block_types",
+    }
+    return required.issubset(config)
+
+
+def _find_model_config(model_dir: Path):
+    """Find the Aniimage model_config.json that describes training settings."""
+    if not model_dir.exists():
+        return None
+    candidates = [
+        p for p in model_dir.rglob("model_config.json")
+        if p.is_file() and ".cache" not in p.parts
+    ]
+    for path in sorted(candidates, key=lambda p: (len(p.relative_to(model_dir).parts), str(p))):
+        config = _read_json(path)
+        if isinstance(config.get("unet"), dict):
+            return path
+    return None
+
+
+def _find_unet_assets(model_dir: Path):
+    """Find Aniimage UNet weights/config, including nested repo folders.
+
+    Aniimage-2 has been published with an extra ``Aniimage-2/unet`` directory
+    inside the repository snapshot. Searching recursively keeps the launcher
+    compatible with that layout as well as normal Diffusers layouts.
+    """
+    if not model_dir.exists():
+        return None
+
+    # Prefer the canonical single-file names, but accept fp16/variant names.
+    weight_candidates = []
+    for pattern in (
+        "diffusion_pytorch_model.safetensors",
+        "diffusion_pytorch_model.bin",
+        "diffusion_pytorch_model*.safetensors",
+        "diffusion_pytorch_model*.bin",
+    ):
+        weight_candidates.extend(model_dir.rglob(pattern))
+
+    # Remove duplicates, metadata files, and anything under the HF cache.
+    unique_weights = []
+    seen = set()
+    for path in weight_candidates:
+        if not path.is_file() or path.name.endswith(".index.json"):
+            continue
+        if ".cache" in path.parts:
+            continue
+        key = str(path.resolve())
+        if key not in seen:
+            seen.add(key)
+            unique_weights.append(path)
+
+    if unique_weights:
+        # Exact canonical filenames first, then the shallowest path.
+        def weight_rank(path: Path):
+            exact = path.name in {
+                "diffusion_pytorch_model.safetensors",
+                "diffusion_pytorch_model.bin",
+            }
+            safe = path.suffix.lower() == ".safetensors"
+            return (not exact, not safe, len(path.relative_to(model_dir).parts), str(path))
+
+        weights_path = sorted(unique_weights, key=weight_rank)[0]
+
+        # Search from the weights folder upward, then across the model root.
+        config_candidates = []
+        current = weights_path.parent
+        while True:
+            config_candidates.extend((current / "config.json", current / "model_config.json"))
+            if current == model_dir or model_dir not in current.parents:
+                break
+            current = current.parent
+        config_candidates.extend(model_dir.rglob("config.json"))
+        config_candidates.extend(model_dir.rglob("model_config.json"))
+
+        config_path = None
+        seen_configs = set()
+        for candidate in config_candidates:
+            if not candidate.is_file() or ".cache" in candidate.parts:
+                continue
+            key = str(candidate.resolve())
+            if key in seen_configs:
+                continue
+            seen_configs.add(key)
+            if _looks_like_unet_config(_read_json(candidate)):
+                config_path = candidate
+                break
+
+        return {
+            "kind": "diffusers_weights",
+            "weights": weights_path,
+            "config": config_path,
+        }
+
+    # Older local checkpoints are still supported.
+    for filename in ("ema_unet.pt", "unet.pt"):
+        candidates = [p for p in model_dir.rglob(filename)
+                      if p.is_file() and ".cache" not in p.parts]
+        if candidates:
+            checkpoint = sorted(
+                candidates,
+                key=lambda p: (len(p.relative_to(model_dir).parts), str(p)),
+            )[0]
+            return {
+                "kind": "checkpoint",
+                "weights": checkpoint,
+                "config": None,
+            }
+
+    return None
+
+
+def _detect_model_resolution(model_dir: Path) -> str:
+    """Infer output resolution from model metadata or the UNet config."""
+    model_config_path = _find_model_config(model_dir)
+    if model_config_path:
+        image_size = _read_json(model_config_path).get("image_size")
+        if isinstance(image_size, int) and image_size > 0:
+            return str(image_size)
+
+    assets = _find_unet_assets(model_dir)
+    config_path = assets.get("config") if assets else None
+    if config_path:
+        config = _read_json(config_path)
+        sample_size = config.get("sample_size")
+        if isinstance(sample_size, (list, tuple)) and sample_size:
+            sample_size = sample_size[0]
+        if isinstance(sample_size, int) and sample_size > 0:
+            return str(sample_size * 8)
+
+    return "512" if "aniimage-2" in model_dir.name.lower() else "256"
+
+
+def download_from_hf():
+    """Download Aniimage-2 from Hugging Face if it is not already present."""
+    try:
+        from huggingface_hub import snapshot_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-2"
+
+    existing = _find_unet_assets(aniimage_dir)
+    existing_config = _find_model_config(aniimage_dir)
+    if existing and existing_config:
+        print(f"Aniimage-2 weights already downloaded: {existing['weights']}")
+        return aniimage_dir
+
+    print(f"Downloading Aniimage-2 from {HF_REPO_ID}...")
+    aniimage_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        snapshot_download(
+            repo_id=HF_REPO_ID,
+            local_dir=aniimage_dir,
+            allow_patterns=[
+                "Aniimage-2/model_config.json",
+                "Aniimage-2/unet/*",
+            ],
+        )
+    except Exception as exc:
+        print(f"Aniimage-2 download failed: {exc}")
+        return None
+
+    assets = _find_unet_assets(aniimage_dir)
+    if not assets:
+        print(
+            "Aniimage-2 repository downloaded, but no supported UNet weights "
+            "were found anywhere below:\n"
+            f"  {aniimage_dir}\n"
+            "Expected diffusion_pytorch_model.safetensors or "
+            "diffusion_pytorch_model.bin."
+        )
+        return None
+
+    print(f"Download complete! Found weights at: {assets['weights']}")
+    return aniimage_dir
+
+
+def find_models():
+    """Find models, including checkpoints nested inside repository folders."""
+    options = []
+    if MODEL_DIR.exists():
+        for d in sorted(MODEL_DIR.iterdir()):
+            if not d.is_dir():
+                continue
+            assets = _find_unet_assets(d)
+            if not assets:
+                continue
+            resolution = _detect_model_resolution(d)
+            model_kind = (
+                "safetensors"
+                if assets["weights"].suffix.lower() == ".safetensors"
+                else assets["kind"]
+            )
+            options.append((model_kind, d.name, d, resolution))
+    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.dtype = self._select_dtype()
+        self.vae = None
+        self.text_encoder = None
+        self.tokenizer = None
+        self.unet = None
+        self.scheduler = None
+        self.loaded_checkpoint = None
+        self.loaded_vae_id = None
+        self.model_config = {}
+        self._clip_inner = None
+        self._clip_full_layers = None
+        self.latent_size = 64
+        self.output_size = 512
+        self.prediction_type = "v_prediction"
+        self.zero_terminal_snr = True
+        self.timestep_spacing = "trailing"
+        self.guidance_rescale = 0.7
+        self.num_train_timesteps = 1000
+        self.beta_schedule = "scaled_linear"
+        self.clip_penultimate = True
+        self.vae_id = VAE_ID
+        self.scheduler_name = "DPM++ SDE Karras"
+        self.cancelled = False
+        self._configure_backends()
+
+    def _select_dtype(self):
+        if self.device != "cuda":
+            return torch.float32
+        bf16_supported = getattr(torch.cuda, "is_bf16_supported", lambda: False)()
+        return torch.bfloat16 if bf16_supported else torch.float16
+
+    def _configure_backends(self):
+        if self.device == "cuda":
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+            if hasattr(torch, "set_float32_matmul_precision"):
+                torch.set_float32_matmul_precision("high")
+
+    def _autocast(self):
+        return torch.autocast(
+            device_type="cuda",
+            dtype=self.dtype,
+            enabled=(self.device == "cuda"),
+        )
+
+    def switch_device(self, new_device):
+        """Switch device and rebuild the models in the correct precision."""
+        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
+        self.dtype = self._select_dtype()
+        self._configure_backends()
+        self.vae = None
+        self.text_encoder = None
+        self.tokenizer = None
+        self.unet = None
+        self.scheduler = None
+        self.loaded_checkpoint = None
+        self.loaded_vae_id = None
+        self._clip_inner = None
+        self._clip_full_layers = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"Switched to {self.device.upper()} ({self.dtype})")
+
+    def _load_model_metadata(self, model_path: Path, res_label: str):
+        """Load the exact training objective and component IDs for Aniimage-2."""
+        config_path = _find_model_config(model_path)
+        config = _read_json(config_path) if config_path else {}
+        self.model_config = config
+
+        self.prediction_type = config.get("prediction_type", "v_prediction")
+        self.zero_terminal_snr = bool(config.get("zero_terminal_snr", True))
+        self.timestep_spacing = config.get(
+            "timestep_spacing",
+            "trailing" if self.zero_terminal_snr else "leading",
+        )
+        self.guidance_rescale = float(config.get("guidance_rescale", 0.7))
+        self.num_train_timesteps = int(config.get("num_train_timesteps", 1000))
+        self.beta_schedule = config.get("beta_schedule", "scaled_linear")
+        self.clip_penultimate = bool(config.get("clip_penultimate", True))
+        self.vae_id = config.get("vae", VAE_ID)
+
+        try:
+            fallback_size = int(res_label)
+        except (TypeError, ValueError):
+            fallback_size = 512
+        self.output_size = int(config.get("image_size", fallback_size))
+        self.latent_size = self.output_size // 8
+
+        if config_path:
+            print(f"Using model metadata: {config_path}")
+        else:
+            print("model_config.json was not found; using Aniimage-2 defaults.")
+
+    def _apply_clip_layer_mode(self):
+        if self.text_encoder is None:
+            return
+        self._clip_inner = getattr(self.text_encoder, "text_model", self.text_encoder)
+        if self._clip_full_layers is None:
+            self._clip_full_layers = self._clip_inner.encoder.layers
+        if self.clip_penultimate:
+            self._clip_inner.encoder.layers = self._clip_full_layers[:-1]
+            print("Text encoder: CLIP penultimate layer (matches training).")
+        else:
+            self._clip_inner.encoder.layers = self._clip_full_layers
+
+    def load_shared(self):
+        from diffusers import AutoencoderKL
+        from transformers import (CLIPConfig, CLIPTextConfig,
+                                  CLIPTextModel, CLIPTokenizer)
+
+        load_kwargs = {"low_cpu_mem_usage": True}
+        if self.device == "cuda":
+            load_kwargs["torch_dtype"] = self.dtype
+
+        if self.vae is None or self.loaded_vae_id != self.vae_id:
+            print(f"Loading VAE: {self.vae_id}...")
+            self.vae = AutoencoderKL.from_pretrained(
+                self.vae_id,
+                **load_kwargs,
+            ).to(self.device).eval()
+            self.vae.requires_grad_(False)
+            self.vae.enable_slicing()
+            if self.device == "cuda":
+                self.vae.to(memory_format=torch.channels_last)
+            self.loaded_vae_id = self.vae_id
+
+        if self.text_encoder is None:
+            print(f"Loading CLIP text encoder: {CLIP_ID}...")
+            self.tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID)
+
+            # Explicitly pass the nested text config. This avoids the
+            # CLIPConfig.hidden_size crash seen with some Transformers builds.
+            clip_config = CLIPConfig.from_pretrained(CLIP_ID)
+            text_config = getattr(clip_config, "text_config", None)
+            if isinstance(text_config, dict):
+                text_config = CLIPTextConfig.from_dict(text_config)
+            if not isinstance(text_config, CLIPTextConfig):
+                text_config = CLIPTextConfig.from_pretrained(CLIP_ID)
+
+            self.text_encoder = CLIPTextModel.from_pretrained(
+                CLIP_ID,
+                config=text_config,
+                **load_kwargs,
+            ).to(self.device).eval()
+            self.text_encoder.requires_grad_(False)
+            self._clip_full_layers = None
+
+        self._apply_clip_layer_mode()
+        self.scheduler = self._make_scheduler(self.scheduler_name)
+        print("Shared models loaded.")
+
+    def _make_scheduler(self, name="DPM++ SDE Karras"):
+        from diffusers import (DDIMScheduler, DPMSolverMultistepScheduler,
+                               EulerAncestralDiscreteScheduler,
+                               EulerDiscreteScheduler)
+        base = dict(
+            num_train_timesteps=self.num_train_timesteps,
+            beta_schedule=self.beta_schedule,
+            prediction_type=self.prediction_type,
+            rescale_betas_zero_snr=self.zero_terminal_snr,
+            timestep_spacing=self.timestep_spacing,
+        )
+        if name == "DPM++ 2M Karras":
+            return DPMSolverMultistepScheduler(
+                **base, algorithm_type="dpmsolver++",
+                solver_order=2, use_karras_sigmas=True)
+        if name == "DPM++ SDE Karras":
+            return DPMSolverMultistepScheduler(
+                **base, algorithm_type="sde-dpmsolver++",
+                solver_order=2, use_karras_sigmas=True)
+        if name == "Euler a":
+            return EulerAncestralDiscreteScheduler(**base)
+        if name == "Euler":
+            return EulerDiscreteScheduler(**base)
+        return DDIMScheduler(
+            **base, clip_sample=False, set_alpha_to_one=False)
+
+    def set_scheduler(self, name):
+        self.scheduler_name = name
+        self.scheduler = self._make_scheduler(name)
+
+    def load_model(self, model_path: Path, res_label: str = "512"):
+        if str(model_path) == self.loaded_checkpoint:
+            return
+        from diffusers import UNet2DConditionModel
+
+        assets = _find_unet_assets(model_path)
+        if not assets:
+            raise FileNotFoundError(
+                f"No supported UNet weights found anywhere inside {model_path}"
+            )
+
+        self._load_model_metadata(model_path, res_label)
+        self.load_shared()
+
+        weights_path = assets["weights"]
+        config_path = assets.get("config")
+        suffix = weights_path.suffix.lower()
+        same_dir_config = weights_path.parent / "config.json"
+
+        print(
+            f"Loading UNet from {weights_path} "
+            f"({self.output_size}px, {self.prediction_type}, {self.dtype})..."
+        )
+
+        self.unet = None
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        loaded_directly = False
+        if same_dir_config.exists() and _looks_like_unet_config(_read_json(same_dir_config)):
+            try:
+                kwargs = {"low_cpu_mem_usage": True}
+                if self.device == "cuda":
+                    kwargs["torch_dtype"] = self.dtype
+                if suffix == ".safetensors":
+                    kwargs["use_safetensors"] = True
+                elif suffix == ".bin":
+                    kwargs["use_safetensors"] = False
+                self.unet = UNet2DConditionModel.from_pretrained(
+                    weights_path.parent,
+                    **kwargs,
+                ).to(self.device)
+                loaded_directly = True
+                print("Loaded the repository UNet config and weights directly.")
+            except Exception as exc:
+                print(f"Direct Diffusers load failed ({exc}); loading manually.")
+
+        if not loaded_directly:
+            if config_path:
+                unet_config = _read_json(config_path)
+                print(f"Using UNet config: {config_path}")
+            elif isinstance(self.model_config.get("unet"), dict):
+                unet_config = dict(self.model_config["unet"])
+                print("Using UNet config from model_config.json.")
+            else:
+                unet_config = dict(UNET_CONFIG)
+                print("Using built-in Aniimage-2 UNet config.")
+            unet_config["sample_size"] = self.latent_size
+            self.unet = UNet2DConditionModel.from_config(unet_config)
+
+            if suffix == ".safetensors":
+                from safetensors.torch import load_file
+                state = load_file(str(weights_path), device="cpu")
+            else:
+                try:
+                    state = torch.load(weights_path, map_location="cpu", weights_only=True)
+                except TypeError:
+                    state = torch.load(weights_path, map_location="cpu")
+
+            if weights_path.name == "ema_unet.pt" and isinstance(state, dict) and "shadow_params" in state:
+                params = dict(self.unet.named_parameters())
+                keys = list(params.keys())
+                if len(state["shadow_params"]) != len(keys):
+                    raise RuntimeError("EMA parameter count does not match the UNet.")
+                for key, shadow_param in zip(keys, state["shadow_params"]):
+                    params[key].data.copy_(shadow_param)
+            else:
+                if isinstance(state, dict) and "state_dict" in state:
+                    state = state["state_dict"]
+                if isinstance(state, dict) and state and all(
+                    isinstance(key, str) and key.startswith("module.") for key in state
+                ):
+                    state = {key[7:]: value for key, value in state.items()}
+                self.unet.load_state_dict(state, strict=True)
+
+            if self.device == "cuda":
+                self.unet = self.unet.to(device=self.device, dtype=self.dtype)
+            else:
+                self.unet = self.unet.to(self.device)
+
+        sample_size = self.unet.config.sample_size
+        if isinstance(sample_size, (list, tuple)) and sample_size:
+            sample_size = sample_size[0]
+        if isinstance(sample_size, int) and sample_size > 0:
+            self.latent_size = sample_size
+            vae_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+            self.output_size = sample_size * vae_factor
+
+        self.unet.eval().requires_grad_(False)
+        if self.device == "cuda":
+            self.unet.to(memory_format=torch.channels_last)
+        self.scheduler = self._make_scheduler(self.scheduler_name)
+        self.loaded_checkpoint = str(model_path)
+        print(
+            f"Ready at {self.output_size}x{self.output_size}; "
+            f"zero-SNR={self.zero_terminal_snr}, spacing={self.timestep_spacing}, "
+            f"CFG rescale={self.guidance_rescale}."
+        )
+
+    def _encode_prompts(self, prompt: str, negative_prompt: str):
+        tokens = self.tokenizer(
+            [negative_prompt or "", prompt],
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        with self._autocast():
+            return self.text_encoder(tokens.input_ids.to(self.device))[0]
+
+    @staticmethod
+    def _cfg_rescale(noise_cfg, noise_text, amount):
+        if amount <= 0:
+            return noise_cfg
+        dims = tuple(range(1, noise_cfg.ndim))
+        std_text = noise_text.std(dim=dims, keepdim=True)
+        std_cfg = noise_cfg.std(dim=dims, keepdim=True).clamp_min(1e-6)
+        noise_rescaled = noise_cfg * (std_text / std_cfg)
+        return amount * noise_rescaled + (1.0 - amount) * noise_cfg
+
+    def _decode_latents(self, latents, post_process=False):
+        del post_process  # Kept for compatibility with the preview callbacks.
+        decode_dtype = self.dtype if self.device == "cuda" else torch.float32
+        scaled = (latents / self.vae.config.scaling_factor).to(dtype=decode_dtype)
+        with self._autocast():
+            image = self.vae.decode(scaled).sample
+        image = (image.float() / 2 + 0.5).clamp(0, 1)
+        image = image[0].cpu().permute(1, 2, 0).numpy()
+        image = (image * 255).round().astype("uint8")
+        return Image.fromarray(image)
+
+    @torch.inference_mode()
+    def generate(self, prompt: str, negative_prompt: str = "",
+                 steps: int = 50, 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()
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        embeddings = self._encode_prompts(prompt, negative_prompt)
+
+        scheduler = self._make_scheduler(self.scheduler_name)
+        scheduler.set_timesteps(int(steps), device=self.device)
+
+        in_channels = int(self.unet.config.in_channels)
+        latents = torch.randn(
+            (1, in_channels, self.latent_size, self.latent_size),
+            generator=generator,
+            device=self.device,
+            dtype=torch.float32,
+        ) * scheduler.init_noise_sigma
+
+        total_steps = len(scheduler.timesteps)
+        preview_interval = max(1, int(preview_every))
+
+        for step_i, timestep in enumerate(scheduler.timesteps):
+            if self.cancelled:
+                return None, seed
+
+            latent_input = torch.cat([latents, latents], dim=0)
+            latent_input = scheduler.scale_model_input(latent_input, timestep)
+
+            with self._autocast():
+                prediction = self.unet(
+                    latent_input,
+                    timestep,
+                    encoder_hidden_states=embeddings,
+                ).sample
+
+            pred_negative, pred_text = prediction.chunk(2)
+            prediction = pred_negative + float(guidance_scale) * (pred_text - pred_negative)
+            prediction = self._cfg_rescale(
+                prediction, pred_text, self.guidance_rescale)
+            latents = scheduler.step(prediction, timestep, latents).prev_sample
+
+            if (preview_callback
+                    and (step_i + 1) % preview_interval == 0
+                    and step_i < total_steps - 1):
+                preview_callback(
+                    self._decode_latents(latents),
+                    step_i + 1,
+                    total_steps,
+                )
+
+        return self._decode_latents(latents), seed
+
+    @torch.inference_mode()
+    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.convert("RGB").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)
+        img_tensor = img_tensor.to(
+            dtype=self.dtype if self.device == "cuda" else torch.float32)
+
+        with self._autocast():
+            latents = self.vae.encode(img_tensor).latent_dist.sample()
+        latents = (latents * self.vae.config.scaling_factor).float()
+        embeddings = self._encode_prompts(prompt, negative_prompt)
+
+        scheduler = self._make_scheduler(self.scheduler_name)
+        scheduler.set_timesteps(int(extra_steps), device=self.device)
+        start_step = max(0, int(len(scheduler.timesteps) * (1.0 - float(strength))))
+        timesteps = scheduler.timesteps[start_step:]
+        if len(timesteps) == 0:
+            return source_image.copy()
+
+        noise = torch.randn_like(latents)
+        latents = scheduler.add_noise(latents, noise, timesteps[:1])
+        total_steps = len(timesteps)
+        preview_interval = max(1, int(preview_every))
+
+        for step_i, timestep in enumerate(timesteps):
+            if self.cancelled:
+                return None
+
+            latent_input = torch.cat([latents, latents], dim=0)
+            latent_input = scheduler.scale_model_input(latent_input, timestep)
+            with self._autocast():
+                prediction = self.unet(
+                    latent_input,
+                    timestep,
+                    encoder_hidden_states=embeddings,
+                ).sample
+
+            pred_negative, pred_text = prediction.chunk(2)
+            prediction = pred_negative + float(guidance_scale) * (pred_text - pred_negative)
+            prediction = self._cfg_rescale(
+                prediction, pred_text, self.guidance_rescale)
+            latents = scheduler.step(prediction, timestep, latents).prev_sample
+
+            if (preview_callback
+                    and (step_i + 1) % preview_interval == 0
+                    and step_i < total_steps - 1):
+                preview_callback(
+                    self._decode_latents(latents),
+                    step_i + 1,
+                    total_steps,
+                )
+
+        return self._decode_latents(latents)
+
+
+# ── 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++ SDE 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="50")
+        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)
+
+        # ── 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()
+
+            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
+
+                    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()
+
+            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
+
+                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()