app.py

"""
DiffuseCraft Mod - Improved Version
====================================
Original: https://huggingface.co/spaces/R-Kentaren/DiffuseCraftMod

IMPROVEMENTS:
- Bug fixes: Error handling, memory leaks, race conditions
- New features: Batch generation, smart presets, prompt templates, enhanced cache
- Optimizations: GPU memory management, thread safety
- NO UI/CSS/THEME CHANGES (as requested)
"""

import spaces
import os
from argparse import ArgumentParser
from stablepy import (
    Model_Diffusers,
    SCHEDULE_TYPE_OPTIONS,
    SCHEDULE_PREDICTION_TYPE_OPTIONS,
    check_scheduler_compatibility,
    TASK_AND_PREPROCESSORS,
    FACE_RESTORATION_MODELS,
    PROMPT_WEIGHT_OPTIONS_PRIORITY,
    scheduler_names,
)
from constants import (
    DIRECTORY_UPSCALERS,
    TASK_STABLEPY,
    TASK_MODEL_LIST,
    UPSCALER_DICT_GUI,
    UPSCALER_KEYS,
    PROMPT_W_OPTIONS,
    WARNING_MSG_VAE,
    SDXL_TASK,
    MODEL_TYPE_TASK,
    POST_PROCESSING_SAMPLER,
    DIFFUSERS_CONTROLNET_MODEL,
    IP_MODELS,
    MODE_IP_OPTIONS,
    CACHE_HF_ROOT,
)
from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
import torch
import re
import time
import threading
from PIL import ImageFile
from utils import (
    get_model_list,
    extract_parameters,
    get_model_type,
    extract_exif_data,
    create_mask_now,
    download_diffuser_repo,
    get_used_storage_gb,
    delete_model,
    progress_step_bar,
    html_template_message,
    escape_html,
    clear_hf_cache,
)
from image_processor import preprocessor_tab
from datetime import datetime
import gradio as gr
import logging
import diffusers
import warnings
from stablepy import logger
from diffusers import FluxPipeline
import subprocess
import json
from pathlib import Path
from typing import Generator, Tuple, List, Any, Optional, Dict
import traceback
import hashlib
import copy
from contextlib import contextmanager
import functools

# ==================== CONFIGURATION ====================
IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
if IS_ZERO_GPU:
    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
IS_GPU_MODE = True if IS_ZERO_GPU else (True if torch.cuda.is_available() else False)
img_path = "./images/"
allowed_path = os.path.abspath(img_path)
delete_cache_time = (9600, 9600) if IS_ZERO_GPU else (86400, 86400)

ImageFile.LOAD_TRUNCATED_IMAGES = True
torch.backends.cuda.matmul.allow_tf32 = True

# ==================== IMPROVED IMPORTS ====================
from modutils import (list_uniq, download_private_repo, get_model_id_list, get_tupled_embed_list,
    get_lora_model_list, get_all_lora_tupled_list, update_loras, apply_lora_prompt, set_prompt_loras,
    get_my_lora, upload_file_lora, move_file_lora, search_civitai_lora, select_civitai_lora,
    update_civitai_selection, get_civitai_tag, CIVITAI_SORT, CIVITAI_PERIOD, CIVITAI_BASEMODEL,
    set_textual_inversion_prompt, get_model_pipeline, change_interface_mode, get_t2i_model_info, download_link_model,
    get_tupled_model_list, save_gallery_images, save_gallery_history, set_optimization, set_sampler_settings,
    set_quick_presets, process_style_prompt, optimization_list, save_images, download_things, valid_model_name,
    preset_styles, preset_quality, preset_sampler_setting, translate_to_en, EXAMPLES_GUI, RESOURCES)
from env import (HF_TOKEN, CIVITAI_API_KEY, HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
    HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
    DIRECTORY_MODELS, DIRECTORY_LORAS, DIRECTORY_VAES, DIRECTORY_EMBEDS, DIRECTORY_EMBEDS_SDXL,
    DIRECTORY_EMBEDS_POSITIVE_SDXL, LOAD_DIFFUSERS_FORMAT_MODEL,
    DOWNLOAD_MODEL_LIST, DOWNLOAD_LORA_LIST, DOWNLOAD_VAE_LIST, DOWNLOAD_EMBEDS)

from tagger.v2 import V2_ALL_MODELS, v2_random_prompt, v2_upsampling_prompt
from tagger.utils import (gradio_copy_text, COPY_ACTION_JS, gradio_copy_prompt,
    V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS, V2_LENGTH_OPTIONS, V2_IDENTITY_OPTIONS)
from tagger.tagger import (predict_tags_wd, convert_danbooru_to_e621_prompt,
    remove_specific_prompt, insert_recom_prompt, insert_model_recom_prompt,
    compose_prompt_to_copy, translate_prompt, select_random_character)

# Download private repos
download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO, DIRECTORY_LORAS, True)
download_private_repo(HF_VAE_PRIVATE_REPO, DIRECTORY_VAES, False)

# Create directories
directories = [DIRECTORY_MODELS, DIRECTORY_LORAS, DIRECTORY_VAES, DIRECTORY_EMBEDS, DIRECTORY_UPSCALERS]
for directory in directories:
    os.makedirs(directory, exist_ok=True)

# Download models/VAEs/LoRAs
DOWNLOAD_MODEL = ", ".join(DOWNLOAD_MODEL_LIST)
DOWNLOAD_VAE = ", ".join(DOWNLOAD_VAE_LIST)
DOWNLOAD_LORA = ", ".join(DOWNLOAD_LORA_LIST)

for url in [url.strip() for url in DOWNLOAD_MODEL.split(',')]:
    download_things(DIRECTORY_MODELS, url, HF_TOKEN, CIVITAI_API_KEY)
for url in [url.strip() for url in DOWNLOAD_VAE.split(',')]:
    download_things(DIRECTORY_VAES, url, HF_TOKEN, CIVITAI_API_KEY)
for url in [url.strip() for url in DOWNLOAD_LORA.split(',')]:
    download_things(DIRECTORY_LORAS, url, HF_TOKEN, CIVITAI_API_KEY)

# Download Embeddings
for url_embed in DOWNLOAD_EMBEDS:
    if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
        download_things(DIRECTORY_EMBEDS, url_embed, HF_TOKEN, CIVITAI_API_KEY)

# Build model lists
embed_list = get_model_list(DIRECTORY_EMBEDS)
lora_model_list = get_lora_model_list()
vae_model_list = get_model_list(DIRECTORY_VAES)
vae_model_list.insert(0, "BakedVAE")
vae_model_list.insert(0, "None")

single_file_model_list = get_model_list(DIRECTORY_MODELS)
model_list = list_uniq(get_model_id_list() + LOAD_DIFFUSERS_FORMAT_MODEL + single_file_model_list)
download_private_repo(HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_SDXL, False)
download_private_repo(HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO, DIRECTORY_EMBEDS_POSITIVE_SDXL, False)
embed_sdxl_list = get_model_list(DIRECTORY_EMBEDS_SDXL) + get_model_list(DIRECTORY_EMBEDS_POSITIVE_SDXL)

def get_embed_list(pipeline_name):
    return get_tupled_embed_list(embed_sdxl_list if pipeline_name == "StableDiffusionXLPipeline" else embed_list)

print('\033[33m🏁 Download and listing of valid models completed.\033[0m')

# ==================== NEW: PRESET MANAGER ====================
class PresetManager:
    """Manages generation presets with save/load functionality."""
    
    PRESET_DIR = Path("presets")
    
    def __init__(self):
        self.PRESET_DIR.mkdir(exist_ok=True)
        self._cache = {}
        self._load_all_presets()
    
    def _load_all_presets(self):
        """Load all presets from disk."""
        try:
            for preset_file in self.PRESET_DIR.glob("*.json"):
                try:
                    data = json.loads(preset_file.read_text(encoding="utf-8"))
                    self._cache[preset_file.stem] = data
                except Exception as e:
                    print(f"[preset] Failed to load {preset_file.name}: {e}")
        except Exception as e:
            print(f"[preset] Failed to load presets: {e}")
    
    def save_preset(self, name: str, params: dict) -> bool:
        """Save a preset to disk."""
        try:
            # Validate name
            if not name or not name.strip():
                raise ValueError("Preset name cannot be empty")
            safe_name = re.sub(r'[^\w\s-]', '', name.strip()).replace(' ', '_')
            
            filepath = self.PRESET_DIR / f"{safe_name}.json"
            filepath.write_text(json.dumps(params, indent=2, ensure_ascii=False), encoding="utf-8")
            self._cache[safe_name] = params
            print(f"[preset] Saved: {safe_name}")
            return True
        except Exception as e:
            print(f"[preset] Save failed: {e}")
            return False
    
    def load_preset(self, name: str) -> Optional[dict]:
        """Load a preset by name."""
        return self._cache.get(name)
    
    def delete_preset(self, name: str) -> bool:
        """Delete a preset."""
        try:
            filepath = self.PRESET_DIR / f"{name}.json"
            if filepath.exists():
                filepath.unlink()
                self._cache.pop(name, None)
                return True
            return False
        except Exception as e:
            print(f"[preset] Delete failed: {e}")
            return False
    
    def list_presets(self) -> List[str]:
        """List all available preset names."""
        return sorted(self._cache.keys())
    
    def export_presets(self, export_path: str) -> bool:
        """Export all presets to a single JSON file."""
        try:
            data = {
                "exported_at": datetime.now().isoformat(),
                "version": "1.0",
                "presets": self._cache
            }
            Path(export_path).write_text(json.dumps(data, indent=2), encoding="utf-8")
            return True
        except Exception as e:
            print(f"[preset] Export failed: {e}")
            return False
    
    def import_presets(self, import_path: str) -> int:
        """Import presets from a JSON file. Returns count of imported presets."""
        try:
            data = json.loads(Path(import_path).read_text(encoding="utf-8"))
            presets = data.get("presets", {})
            count = 0
            for name, params in presets.items():
                if self.save_preset(name, params):
                    count += 1
            return count
        except Exception as e:
            print(f"[preset] Import failed: {e}")
            return 0


# Initialize global preset manager
preset_manager = PresetManager()

# ==================== NEW: PROMPT TEMPLATE SYSTEM ====================
class PromptTemplateSystem:
    """Manages prompt templates with variable substitution."""
    
    TEMPLATES = {
        "anime_basic": {
            "name": "Basic Anime",
            "template": "1girl, solo, {subject}, {quality_tags}, {style_tags}",
            "variables": {
                "subject": "main character/subject description",
                "quality_tags": "masterpiece, best quality",
                "style_tags": "anime style, detailed"
            },
            "description": "Basic template for anime-style generation"
        },
        "portrait": {
            "name": "Portrait",
            "template": "portrait of {subject}, {lighting}, {quality_tags}, detailed face, professional photography",
            "variables": {
                "subject": "person description",
                "lighting": "studio lighting, soft light",
                "quality_tags": "highly detailed, 8k"
            },
            "description": "Professional portrait template"
        },
        "landscape": {
            "name": "Landscape",
            "template": "scenic view, {setting}, {time_of_day}, {weather}, {quality_tags}, atmospheric",
            "variables": {
                "setting": "mountains/ocean/forest etc.",
                "time_of_day": "sunset/sunrise/noon",
                "weather": "clear sky/cloudy/misty",
                "quality_tags": "highly detailed, photorealistic"
            },
            "description": "Scenic landscape template"
        },
        "character_design": {
            "name": "Character Design",
            "template": "character design, {gender}, {hair} hair, {eyes} eyes, {outfit}, {pose}, {expression}, {quality_tags}, white background, reference sheet",
            "variables": {
                "gender": "girl/boy/person",
                "hair": "long/short/color",
                "eyes": "color/style",
                "outfit": "clothing description",
                "pose": "standing/sitting/action pose",
                "expression": "happy/serious/calm"
            },
            "description": "Character design reference template"
        }
    }
    
    @classmethod
    def get_template(cls, template_id: str) -> Optional[dict]:
        """Get template by ID."""
        return cls.TEMPLATES.get(template_id)
    
    @classmethod
    def list_templates(cls) -> List[tuple]:
        """List all templates as (id, name) tuples."""
        return [(tid, t["name"]) for tid, t in cls.TEMPLATES.items()]
    
    @classmethod
    def render_template(cls, template_id: str, variables: dict) -> Optional[str]:
        """Render a template with provided variables."""
        template_data = cls.TEMPLATES.get(template_id)
        if not template_data:
            return None
        
        template = template_data["template"]
        
        # Merge with defaults
        default_vars = template_data.get("variables", {})
        merged_vars = {}
        for var_name, default_desc in default_vars.items():
            merged_vars[var_name] = variables.get(var_name, f"[{default_desc}]")
        
        # Add any extra variables
        merged_vars.update(variables)
        
        # Substitute variables
        try:
            result = template.format(**merged_vars)
            return result
        except KeyError as e:
            print(f"[template] Missing variable: {e}")
            return None
    
    @classmethod
    def add_custom_template(cls, template_id: str, name: str, template: str, 
                           variables: dict, description: str = "") -> bool:
        """Add a custom template."""
        if template_id in cls.TEMPLATES:
            return False
        
        cls.TEMPLATES[template_id] = {
            "name": name,
            "template": template,
            "variables": variables,
            "description": description
        }
        return True


# ==================== NEW: BATCH GENERATION SYSTEM ====================
class BatchGenerator:
    """Handles batch generation with prompt variations."""
    
    MAX_BATCH_SIZE = 20  # Limit batch size to prevent abuse
    
    @staticmethod
    def generate_variations(base_prompt: str, variation_mode: str = "sequential",
                           count: int = 4, seed_start: int = -1) -> List[Tuple[str, int]]:
        """
        Generate prompt variations for batch processing.
        
        Args:
            base_prompt: Base prompt to vary
            variation_mode: 'sequential' (seed variation), 'prompt_permutation', 'aspect_variations'
            count: Number of variations to generate
            seed_start: Starting seed (-1 for random)
        
        Returns:
            List of (modified_prompt, seed) tuples
        """
        variations = []
        
        if variation_mode == "sequential":
            # Simple seed-based variations
            base_seed = seed_start if seed_start > 0 else int(time.time()) % (2**32 - 1)
            for i in range(min(count, BatchGenerator.MAX_BATCH_SIZE)):
                variations.append((base_prompt, base_seed + i))
                
        elif variation_mode == "prompt_permutation":
            # Generate slight prompt modifications
            modifiers = [
                "best quality, masterpiece",
                "highly detailed",
                "professional, sharp focus",
                "cinematic lighting",
                "trending on artstation",
                "digital art, vibrant colors",
                "detailed background",
                "soft lighting, atmospheric"
            ]
            base_seed = seed_start if seed_start > 0 else int(time.time()) % (2**32 - 1)
            for i in range(min(count, BatchGenerator.MAX_BATCH_SIZE)):
                modifier = modifiers[i % len(modifiers)]
                modified_prompt = f"{base_prompt}, {modifier}" if modifier else base_prompt
                variations.append((modified_prompt, base_seed + i))
                
        elif variation_mode == "aspect_variations":
            # Different aspect ratios embedded in prompt
            aspects = [
                ("portrait, tall composition", (832, 1216)),
                ("landscape, wide composition", (1216, 832)),
                ("square composition", (1024, 1024)),
                ("cinematic widescreen", (1280, 720)),
            ]
            base_seed = seed_start if seed_start > 0 else int(time.time()) % (2**32 - 1)
            for i in range(min(count, BatchGenerator.MAX_BATCH_SIZE)):
                aspect_mod, _ = aspects[i % len(aspects)]
                modified_prompt = f"{base_prompt}, {aspect_mod}"
                variations.append((modified_prompt, base_seed + i))
        
        return variations
    
    @staticmethod
    def validate_batch_params(prompt: str, count: int, mode: str) -> Tuple[bool, str]:
        """Validate batch generation parameters."""
        if not prompt or not prompt.strip():
            return False, "Prompt cannot be empty"
        
        if count < 1 or count > BatchGenerator.MAX_BATCH_SIZE:
            return False, f"Batch size must be between 1 and {BatchGenerator.MAX_BATCH_SIZE}"
        
        valid_modes = ["sequential", "prompt_permutation", "aspect_variations"]
        if mode not in valid_modes:
            return False, f"Invalid mode. Must be one of: {valid_modes}"
        
        return True, ""


# ==================== NEW: ENHANCED CACHE MANAGER ====================
class EnhancedCacheManager:
    """Improved cache management with LRU eviction and memory tracking."""
    
    def __init__(self, max_size_gb: float = 4.0, max_files: int = 256):
        self.max_size_bytes = max_size_gb * 1024**3
        self.max_files = max_files
        self.cache_dir = Path("outputs")
        self.cache_dir.mkdir(exist_ok=True)
        self._access_log: Dict[str, float] = {}  # path -> last access time
        self._lock = threading.Lock()
    
    def record_access(self, filepath: str):
        """Record file access for LRU tracking."""
        with self._lock:
            self._access_log[filepath] = time.time()
    
    def get_cache_stats(self) -> dict:
        """Get current cache statistics."""
        try:
            files = list(self.cache_dir.glob("*.png"))
            total_size = sum(f.stat().st_size for f in files if f.is_file())
            return {
                "file_count": len(files),
                "total_size_bytes": total_size,
                "total_size_gb": round(total_size / (1024**3), 3),
                "max_size_gb": round(self.max_size_bytes / (1024**3), 3),
                "usage_percent": round(total_size / self.max_size_bytes * 100, 1) if self.max_size_bytes > 0 else 0
            }
        except Exception as e:
            print(f"[cache] Stats error: {e}")
            return {"error": str(e)}
    
    def cleanup(self, protect_paths: List[str] = None) -> int:
        """
        Clean up cache using LU strategy.
        Returns number of files removed.
        """
        protect_set = {str(Path(p).resolve()) for p in (protect_paths or [])}
        removed_count = 0
        
        with self._lock:
            try:
                files = []
                for f in self.cache_dir.glob("*.png"):
                    if f.is_file():
                        try:
                            st = f.stat()
                            last_access = self._access_log.get(str(f), st.st_mtime)
                            files.append((f, st.st_size, last_access))
                        except Exception:
                            continue
                
                # Sort by last access time (oldest first)
                files.sort(key=lambda x: x[2])
                
                total_size = sum(size for _, size, _ in files)
                
                # Remove files until under limits
                while (files and 
                      (len(files) > self.max_files or total_size > self.max_size_bytes)):
                    fpath, fsize, _ = files.pop(0)
                    
                    if str(fpath.resolve()) in protect_set:
                        continue
                    
                    try:
                        fpath.unlink()
                        total_size -= fsize
                        self._access_log.pop(str(fpath), None)
                        removed_count += 1
                        print(f"[cache] Removed: {fpath.name}")
                    except Exception as e:
                        print(f"[cache] Remove failed {fpath.name}: {e}")
                        
            except Exception as e:
                print(f"[cache] Cleanup error: {e}")
        
        return removed_count
    
    def clear_old_entries(self, max_age_seconds: float = 86400) -> int:
        """Remove entries older than max_age_seconds."""
        cutoff = time.time() - max_age_seconds
        removed = 0
        
        with self._lock:
            try:
                for fpath, access_time in list(self._access_log.items()):
                    if access_time < cutoff:
                        p = Path(fpath)
                        if p.exists():
                            try:
                                p.unlink()
                                removed += 1
                            except Exception:
                                pass
                        del self._access_log[fpath]
            except Exception as e:
                print(f"[cache] Clear old error: {e}")
        
        return removed


# Initialize enhanced cache manager
cache_manager = EnhancedCacheManager()


# ==================== IMPROVED GuiSD CLASS ====================
components = None
if IS_ZERO_GPU:
    flux_repo = "camenduru/FLUX.1-dev-diffusers"
    flux_pipe = FluxPipeline.from_pretrained(
        flux_repo,
        transformer=None,
        torch_dtype=torch.bfloat16,
    )
    components = flux_pipe.components
    delete_model(flux_repo)

#######################
# GUI
#######################
logging.getLogger("diffusers").setLevel(logging.ERROR)
diffusers.utils.logging.set_verbosity(40)
warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffusers")
warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
warnings.filterwarnings(action="ignore", category=FutureWarning, module="transformers")

parser = ArgumentParser(description='DiffuseCraft: Create images from text prompts.', add_help=True)
parser.add_argument("--share", action="store_true", dest="share_enabled", default=False, help="Enable sharing")
parser.add_argument('--theme', type=str, default="NoCrypt/miku", help='Set the theme (default: NoCrypt/miku)')
parser.add_argument("--ssr", action="store_true", default=False, help="Enable SSR (Server-Side Rendering)")
parser.add_argument("--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set logging level (default: INFO)")
args = parser.parse_args()

logger.setLevel(
    "INFO" if IS_ZERO_GPU else getattr(logging, args.log_level.upper())
)


def description_ui():
    gr.Markdown(
        """
## Danbooru Tags Transformer V2 Demo with WD Tagger
(Image =>) Prompt => Upsampled longer prompt
- Mod of p1atdev's [Danbooru Tags Transformer V2 Demo](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) and [WD Tagger with 🤗 transformers](https://huggingface.co/spaces/p1atdev/wd-tagger-transformers).
- Models: p1atdev's [wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf), [dart-v2-moe-sft](https://huggingface.co/p1atdev/dart-v2-moe-sft)
"""
    )


def lora_chk(lora_):
    if isinstance(lora_, str) and lora_.strip() not in ["", "None"]:
        return lora_
    return None


# Context manager for GPU operations
@contextmanager
def gpu_context(duration: int = 60):
    """Context manager for GPU operations with automatic cleanup."""
    try:
        if IS_ZERO_GPU:
            yield spaces.GPU(duration=duration)
        else:
            yield None
    finally:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


class GuiSD:
    """Improved GUI SD class with better error handling and resource management."""
    
    def __init__(self, stream=True):
        self.model = None
        self.status_loading = False
        self.sleep_loading = 4
        self.last_load = datetime.now()
        self.inventory = []

        # Avoid duplicate downloads - FIXED: Better thread safety
        self.active_downloads = set()
        self.download_lock = threading.RLock()  # Changed to RLock for reentrant locking
        self.download_events = {}  # For per-model wait events

        # Anti-abuse: track new model requests.
        self.used_models = []
        self.new_model_history = []
        
        # Generation statistics
        self.generation_stats = {
            "total_generations": 0,
            "successful_generations": 0,
            "failed_generations": 0,
            "last_error": None
        }

    def update_storage_models(self, storage_floor_gb=24, required_inventory_for_purge=3):
        """Update storage and purge old models if needed."""
        try:
            while get_used_storage_gb() > storage_floor_gb:
                if len(self.inventory) < required_inventory_for_purge:
                    break
                removal_candidate = self.inventory.pop(0)
                delete_model(removal_candidate)

            # Cleanup after 60 seconds of inactivity
            lowPrioCleanup = max((datetime.now() - self.last_load).total_seconds(), 0) > 120
            if lowPrioCleanup and not self.status_loading and get_used_storage_gb(CACHE_HF_ROOT) > (storage_floor_gb * 2):
                print("Cleaning up Hugging Face cache...")
                clear_hf_cache()
                self.inventory = [
                    m for m in self.inventory if os.path.exists(m)
                ]
        except Exception as e:
            print(f"[storage] Update error: {e}")

    def update_inventory(self, model_name):
        """Update model inventory with proper error handling."""
        try:
            if model_name not in single_file_model_list:
                self.inventory = [
                    m for m in self.inventory if m != model_name
                ] + [model_name]
            print(self.inventory)
        except Exception as e:
            print(f"[inventory] Update error: {e}")

    def load_new_model(self, model_name, vae_model, task, controlnet_model, progress=gr.Progress(track_tqdm=True)):
        """
        Load new model with improved error handling and resource management.
        
        FIXES:
        - Better exception handling
        - Proper resource cleanup in finally blocks
        - Improved thread safety
        - Memory leak prevention
        """
        loaded_successfully = False
        
        try:
            if model_name != model_list[0]:
                # --- Anti-Abuse Check Start ---
                if model_name in self.used_models:
                    # Move to the end to mark as the most recently used.
                    self.used_models.remove(model_name)
                    self.used_models.append(model_name)
                else:
                    current_time = datetime.now()
                    # Retain history of new model requests from the last 20 minutes.
                    self.new_model_history = [
                        t for t in self.new_model_history
                        if (current_time - t).total_seconds() < 1200
                    ]

                    # Allow a maximum of 5 new model requests per 20 minutes.
                    if len(self.new_model_history) >= 5:
                        yield "Rate limit exceeded: Too many new models requested."
                        raise gr.Error("Too many new models requested. Please reuse your previously loaded models or wait a few minutes before trying new ones.")

                    self.new_model_history.append(current_time)
                    self.used_models.append(model_name)

                    # Cap the reuse list to the 5 most recent models.
                    if len(self.used_models) > 5:
                        self.used_models.pop(0)
                # --- Anti-Abuse Check End ---

            lock_key = model_name
            
            # Improved waiting mechanism with timeout
            wait_timeout = 300  # 5 minutes max wait
            wait_start = time.time()
            
            while True:
                with self.download_lock:
                    if lock_key not in self.active_downloads:
                        self.active_downloads.add(lock_key)
                        break
                
                # Check timeout
                elapsed = time.time() - wait_start
                if elapsed > wait_timeout:
                    yield f"Timeout waiting for download: {model_name}"
                    raise TimeoutError(f"Waited too long for model download: {model_name}")
                
                yield f"Waiting for existing download to finish: {model_name}... ({int(elapsed)}s)"
                time.sleep(1)

            try:
                # Download link model > model_name
                is_link_model = False
                model_type = None
                
                if model_name.startswith("http"):
                    yield f"Downloading model: {model_name}"
                    model_name, model_type = download_link_model(model_name, DIRECTORY_MODELS)
                    if not model_name:
                        raise ValueError("Error retrieving model information from URL")
                    is_link_model = True
                else:
                    is_link_model = False

                if IS_ZERO_GPU:
                    self.update_storage_models()

                vae_model = vae_model if vae_model != "None" else None
                model_type = get_model_type(model_name) if not is_link_model else model_type
                dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16

                if not os.path.exists(model_name):
                    logger.debug(f"model_name={model_name}, vae_model={vae_model}, task={task}, controlnet_model={controlnet_model}")
                    _ = download_diffuser_repo(
                        repo_name=model_name,
                        model_type=model_type,
                        revision="main",
                        token=True,
                    )

                self.update_inventory(model_name)

            finally:
                with self.download_lock:
                    self.active_downloads.discard(lock_key)
            
            # Improved queue waiting with better feedback
            max_queue_wait = 120  # 2 minutes max queue wait
            queue_wait_start = time.time()
            
            for i in range(max_queue_wait * 2):  # Check every 0.5 seconds
                if not self.status_loading:
                    self.status_loading = True
                    if i > 0:
                        time.sleep(self.sleep_loading)
                        print("Previous model ops...")
                    break
                    
                # Check queue timeout
                elapsed = time.time() - queue_wait_start
                if elapsed > max_queue_wait:
                    yield "Timeout waiting for model queue"
                    raise TimeoutError("Model queue timeout")
                
                time.sleep(0.5)
                print(f"Waiting queue {i}")
                yield f"Waiting in queue... ({int(elapsed)}s)"

            self.status_loading = True

            yield f"Loading model: {model_name}"

            if vae_model == "BakedVAE":
                vae_model = model_name
            elif vae_model:
                vae_type = "SDXL" if "sdxl" in vae_model.lower() else "SD 1.5"
                if model_type != vae_type:
                    gr.Warning(WARNING_MSG_VAE)

            print("Loading model...")

            start_time = time.time()

            try:
                if self.model is None:
                    self.model = Model_Diffusers(
                        base_model_id=model_name,
                        task_name=TASK_STABLEPY[task],
                        vae_model=vae_model,
                        type_model_precision=dtype_model,
                        retain_task_model_in_cache=False,
                        controlnet_model=controlnet_model,
                        device="cpu" if IS_ZERO_GPU else None,
                        env_components=components,
                    )
                    self.model.advanced_params(image_preprocessor_cuda_active=IS_GPU_MODE)
                else:
                    if self.model.base_model_id != model_name:
                        load_now_time = datetime.now()
                        elapsed_time = max((load_now_time - self.last_load).total_seconds(), 0)

                        if elapsed_time <= 9:
                            print("Waiting for the previous model's time ops...")
                            time.sleep(9 - elapsed_time)

                    if IS_ZERO_GPU:
                        self.model.device = torch.device("cpu")
                    self.model.load_pipe(
                        model_name,
                        task_name=TASK_STABLEPY[task],
                        vae_model=vae_model,
                        type_model_precision=dtype_model,
                        retain_task_model_in_cache=False,
                        controlnet_model=controlnet_model,
                    )

                end_time = time.time()
                self.sleep_loading = max(min(int(end_time - start_time), 10), 4)
                loaded_successfully = True
                
            except Exception as e:
                # Reset state on error
                self.last_load = datetime.now()
                self.status_loading = False
                self.sleep_loading = 4
                self.generation_stats["last_error"] = str(e)
                raise e

            self.last_load = datetime.now()
            self.status_loading = False

            yield f"Model loaded: {model_name}"
            
        except Exception as e:
            # Ensure status is reset on any error
            self.status_loading = False
            print(f"[model] Load error: {traceback.format_exc()}")
            raise

    #@spaces.GPU
    @torch.inference_mode()
    def generate_pipeline(
        self,
        prompt,
        neg_prompt,
        num_images,
        steps,
        cfg,
        clip_skip,
        seed,
        lora1,
        lora_scale1,
        lora2,
        lora_scale2,
        lora3,
        lora_scale3,
        lora4,
        lora_scale4,
        lora5,
        lora_scale5,
        lora6,
        lora_scale6,
        lora7,
        lora_scale7,
        sampler,
        schedule_type,
        schedule_prediction_type,
        img_height,
        img_width,
        model_name,
        vae_model,
        task,
        image_control,
        preprocessor_name,
        preprocess_resolution,
        image_resolution,
        style_prompt,
        style_json_file,
        image_mask,
        strength,
        low_threshold,
        high_threshold,
        value_threshold,
        distance_threshold,
        recolor_gamma_correction,
        tile_blur_sigma,
        controlnet_output_scaling_in_unet,
        controlnet_start_threshold,
        controlnet_stop_threshold,
        textual_inversion,
        syntax_weights,
        upscaler_model_path,
        upscaler_increases_size,
        upscaler_tile_size,
        upscaler_tile_overlap,
        hires_steps,
        hires_denoising_strength,
        hires_sampler,
        hires_prompt,
        hires_negative_prompt,
        hires_before_adetailer,
        hires_after_adetailer,
        hires_schedule_type,
        hires_guidance_scale,
        controlnet_model,
        loop_generation,
        leave_progress_bar,
        disable_progress_bar,
        image_previews,
        display_images,
        save_generated_images,
        filename_pattern,
        image_storage_location,
        retain_compel_previous_load,
        retain_detailfix_model_previous_load,
        retain_hires_model_previous_load,
        t2i_adapter_preprocessor,
        t2i_adapter_conditioning_scale,
        t2i_adapter_conditioning_factor,
        enable_live_preview,
        freeu,
        generator_in_cpu,
        adetailer_inpaint_only,
        adetailer_verbose,
        adetailer_sampler,
        adetailer_active_a,
        prompt_ad_a,
        negative_prompt_ad_a,
        strength_ad_a,
        face_detector_ad_a,
        person_detector_ad_a,
        hand_detector_ad_a,
        mask_dilation_a,
        mask_blur_a,
        mask_padding_a,
        adetailer_active_b,
        prompt_ad_b,
        negative_prompt_ad_b,
        strength_ad_b,
        face_detector_ad_b,
        person_detector_ad_b,
        hand_detector_ad_b,
        mask_dilation_b,
        mask_blur_b,
        mask_padding_b,
        retain_task_cache_gui,
        guidance_rescale,
        image_ip1,
        mask_ip1,
        model_ip1,
        mode_ip1,
        scale_ip1,
        image_ip2,
        mask_ip2,
        model_ip2,
        mode_ip2,
        scale_ip2,
        pag_scale,
        face_restoration_model,
        face_restoration_visibility,
        face_restoration_weight,
    ):
        """
        Generate images with improved error handling and memory management.
        
        FIXES:
        - Better exception handling throughout
        - Memory leak prevention
        - Proper resource cleanup
        - Generation statistics tracking
        """
        self.generation_stats["total_generations"] += 1
        info_state = html_template_message("Navigating latent space...")
        yield info_state, gr.update(), gr.update()

        try:
            vae_model = vae_model if vae_model != "None" else None
            loras_list = [lora1, lora2, lora3, lora4, lora5, lora6, lora7]
            vae_msg = f"VAE: {vae_model}" if vae_model else ""
            msg_lora = ""

            # FIX: Don't reassign global variable, use local instead
            current_lora_list = get_lora_model_list()
            loras_list = [s if s else "None" for s in loras_list]
            
            lora1, lora_scale1, lora2, lora_scale2, lora3, lora_scale3, lora4, lora_scale4, lora5, lora_scale5, lora6, lora_scale6, lora7, lora_scale7 = \
                    set_prompt_loras(prompt, syntax_weights, model_name, lora1, lora_scale1, lora2, lora_scale2, lora3,
                                      lora_scale3, lora4, lora_scale4, lora5, lora_scale5, lora6, lora_scale6, lora7, lora_scale7)

            logger.debug(f"Config model: {model_name}, {vae_model}, {loras_list}")

            task = TASK_STABLEPY[task]

            params_ip_img = []
            params_ip_msk = []
            params_ip_model = []
            params_ip_mode = []
            params_ip_scale = []

            all_adapters = [
                (image_ip1, mask_ip1, model_ip1, mode_ip1, scale_ip1),
                (image_ip2, mask_ip2, model_ip2, mode_ip2, scale_ip2),
            ]

            if not hasattr(self.model.pipe, "transformer"):
                for imgip, mskip, modelip, modeip, scaleip in all_adapters:
                    if imgip:
                        params_ip_img.append(imgip)
                        if mskip:
                            params_ip_msk.append(mskip)
                        params_ip_model.append(modelip)
                        params_ip_mode.append(modeip)
                        params_ip_scale.append(scaleip)

            concurrency = 5
            self.model.stream_config(concurrency=concurrency, latent_resize_by=1, vae_decoding=False)

            if task != "txt2img" and not image_control:
                raise ValueError("Reference image is required. Please upload one in 'Image ControlNet/Inpaint/Img2img'.")

            if task in ["inpaint", "repaint"] and not image_mask:
                raise ValueError("Mask image not found. Upload one in 'Image Mask' to proceed.")

            if "https://" not in str(UPSCALER_DICT_GUI[upscaler_model_path]):
                upscaler_model = upscaler_model_path
            else:
                url_upscaler = UPSCALER_DICT_GUI[upscaler_model_path]

                if not os.path.exists(f"./{DIRECTORY_UPSCALERS}/{url_upscaler.split('/')[-1]}"):
                    download_things(DIRECTORY_UPSCALERS, url_upscaler, HF_TOKEN)

                upscaler_model = f"./{DIRECTORY_UPSCALERS}/{url_upscaler.split('/')[-1]}"

            logging.getLogger("ultralytics").setLevel(logging.INFO if adetailer_verbose else logging.ERROR)

            adetailer_params_A = {
                "face_detector_ad": face_detector_ad_a,
                "person_detector_ad": person_detector_ad_a,
                "hand_detector_ad": hand_detector_ad_a,
                "prompt": prompt_ad_a,
                "negative_prompt": negative_prompt_ad_a,
                "strength": strength_ad_a,
                "mask_dilation": mask_dilation_a,
                "mask_blur": mask_blur_a,
                "mask_padding": mask_padding_a,
                "inpaint_only": adetailer_inpaint_only,
                "sampler": adetailer_sampler,
            }

            adetailer_params_B = {
                "face_detector_ad": face_detector_ad_b,
                "person_detector_ad": person_detector_ad_b,
                "hand_detector_ad": hand_detector_ad_b,
                "prompt": prompt_ad_b,
                "negative_prompt": negative_prompt_ad_b,
                "strength": strength_ad_b,
                "mask_dilation": mask_dilation_b,
                "mask_blur": mask_blur_b,
                "mask_padding": mask_padding_b,
            }
            
            pipe_params = {
                "prompt": prompt,
                "negative_prompt": neg_prompt,
                "img_height": img_height,
                "img_width": img_width,
                "num_images": num_images,
                "num_steps": steps,
                "guidance_scale": cfg,
                "clip_skip": clip_skip,
                "pag_scale": float(pag_scale),
                "seed": seed,
                "image": image_control,
                "preprocessor_name": preprocessor_name,
                "preprocess_resolution": preprocess_resolution,
                "image_resolution": image_resolution,
                "style_prompt": style_prompt if style_prompt else "",
                "style_json_file": "",
                "image_mask": image_mask,
                "strength": strength,
                "low_threshold": low_threshold,
                "high_threshold": high_threshold,
                "value_threshold": value_threshold,
                "distance_threshold": distance_threshold,
                "recolor_gamma_correction": float(recolor_gamma_correction),
                "tile_blur_sigma": int(tile_blur_sigma),
                "lora_A": lora_chk(lora1),
                "lora_scale_A": lora_scale1,
                "lora_B": lora_chk(lora2),
                "lora_scale_B": lora_scale2,
                "lora_C": lora_chk(lora3),
                "lora_scale_C": lora_scale3,
                "lora_D": lora_chk(lora4),
                "lora_scale_D": lora_scale4,
                "lora_E": lora_chk(lora5),
                "lora_scale_E": lora_scale5,
                "lora_F": lora_chk(lora6),
                "lora_scale_F": lora_scale6,
                "lora_G": lora_chk(lora7),
                "lora_scale_G": lora_scale7,
                "textual_inversion": get_embed_list(self.model.class_name) if textual_inversion else [],
                "syntax_weights": syntax_weights,
                "sampler": sampler,
                "schedule_type": schedule_type,
                "schedule_prediction_type": schedule_prediction_type,
                "xformers_memory_efficient_attention": False,
                "gui_active": True,
                "loop_generation": loop_generation,
                "controlnet_conditioning_scale": float(controlnet_output_scaling_in_unet),
                "control_guidance_start": float(controlnet_start_threshold),
                "control_guidance_end": float(controlnet_stop_threshold),
                "generator_in_cpu": generator_in_cpu,
                "FreeU": freeu,
                "adetailer_A": adetailer_active_a,
                "adetailer_A_params": adetailer_params_A,
                "adetailer_B": adetailer_active_b,
                "adetailer_B_params": adetailer_params_B,
                "leave_progress_bar": leave_progress_bar,
                "disable_progress_bar": disable_progress_bar,
                "image_previews": image_previews,
                "display_images": False,
                "save_generated_images": save_generated_images,
                "filename_pattern": filename_pattern,
                "image_storage_location": image_storage_location,
                "retain_compel_previous_load": retain_compel_previous_load,
                "retain_detailfix_model_previous_load": retain_detailfix_model_previous_load,
                "retain_hires_model_previous_load": retain_hires_model_previous_load,
                "t2i_adapter_preprocessor": t2i_adapter_preprocessor,
                "t2i_adapter_conditioning_scale": float(t2i_adapter_conditioning_scale),
                "t2i_adapter_conditioning_factor": float(t2i_adapter_conditioning_factor),
                "upscaler_model_path": upscaler_model,
                "upscaler_increases_size": upscaler_increases_size,
                "upscaler_tile_size": upscaler_tile_size,
                "upscaler_tile_overlap": upscaler_tile_overlap,
                "hires_steps": hires_steps,
                "hires_denoising_strength": hires_denoising_strength,
                "hires_prompt": hires_prompt,
                "hires_negative_prompt": hires_negative_prompt,
                "hires_sampler": hires_sampler,
                "hires_before_adetailer": hires_before_adetailer,
                "hires_after_adetailer": hires_after_adetailer,
                "hires_schedule_type": hires_schedule_type,
                "hires_guidance_scale": hires_guidance_scale,
                "ip_adapter_image": params_ip_img,
                "ip_adapter_mask": params_ip_msk,
                "ip_adapter_model": params_ip_model,
                "ip_adapter_mode": params_ip_mode,
                "ip_adapter_scale": params_ip_scale,
                "face_restoration_model": face_restoration_model,
                "face_restoration_visibility": face_restoration_visibility,
                "face_restoration_weight": face_restoration_weight,
            }

            # kwargs for diffusers pipeline
            if guidance_rescale:
                pipe_params["guidance_rescale"] = guidance_rescale
            if IS_ZERO_GPU:
                self.model.device = torch.device("cuda:0")
                if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * self.model.num_loras:
                    self.model.pipe.transformer.to(self.model.device)
                    logger.debug("transformer to cuda")

            actual_progress = 0
            info_images = gr.update()
            
            for img, [seed, image_path, metadata] in self.model(**pipe_params):
                info_state = progress_step_bar(actual_progress, steps)
                actual_progress += concurrency
                if image_path:
                    info_images = f"Seeds: {str(seed)}"
                    if vae_msg:
                        info_images = info_images + "<br>" + vae_msg

                    if "Cannot copy out of meta tensor; no data!" in self.model.last_lora_error:
                        msg_ram = "Unable to process the LoRAs due to high RAM usage; please try again later."
                        print(msg_ram)
                        msg_lora += f"<br>{msg_ram}"

                    for status, lora in zip(self.model.lora_status, self.model.lora_memory):
                        if status:
                            msg_lora += f"<br>Loaded: {lora}"
                        elif status is not None:
                            msg_lora += f"<br>Error with: {lora}"

                    if msg_lora:
                        info_images += msg_lora

                    info_images = info_images + "<br>" + "GENERATION DATA:<br>" + escape_html(metadata[-1]) + "<br>-------<br>"

                    download_links = "<br>".join(
                        [
                            f'<a href="{path.replace("/images/", f"/gradio_api/file={allowed_path}/")}" download="{os.path.basename(path)}">Download Image {i + 1}</a>'
                            for i, path in enumerate(image_path)
                        ]
                    )
                    if save_generated_images:
                        info_images += f"<br>{download_links}"
                    
                    img = save_images(img, metadata)
                    
                    # Record in cache manager
                    if image_path:
                        for path in image_path:
                            cache_manager.record_access(path)
                    
                    if not display_images:
                        img = img if img else gr.update()
                    
                    info_state = "COMPLETE"
                    self.generation_stats["successful_generations"] += 1

                elif not enable_live_preview:
                    img = gr.update()

                yield info_state, img, info_images

        except Exception as e:
            self.generation_stats["failed_generations"] += 1
            self.generation_stats["last_error"] = str(e)
            print(f"[generation] Error: {traceback.format_exc()}")
            raise
        finally:
            # Always clean up GPU memory
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()


def dynamic_gpu_duration(func, duration, *args):

    @spaces.GPU(duration=duration)
    def wrapped_func():
        yield from func(*args)

    return wrapped_func()


@spaces.GPU
def dummy_gpu():
    return None


def sd_gen_generate_pipeline(*args):
    """Wrapper for generation pipeline with improved error handling."""
    gpu_duration_arg = int(args[-1]) if args[-1] else 59
    verbose_arg = int(args[-2])
    load_lora_cpu = args[-3]
    generation_args = args[:-3]
    lora_list = [
        None if item == "None" or item == "" else item
        for item in [args[7], args[9], args[11], args[13], args[15], args[17], args[19]]
    ]
    lora_status = [None] * sd_gen.model.num_loras if getattr(sd_gen, "model", None) is not None else 7

    msg_load_lora = "Updating LoRAs in GPU..."
    if load_lora_cpu:
        msg_load_lora = "Updating LoRAs in CPU..."

    if lora_list != sd_gen.model.lora_memory and lora_list != [None] * sd_gen.model.num_loras:
        yield msg_load_lora, gr.update(), gr.update()

    # Load lora in CPU
    if load_lora_cpu:
        try:
            lora_status = sd_gen.model.load_lora_on_the_fly(
                lora_A=lora_list[0], lora_scale_A=args[8],
                lora_B=lora_list[1], lora_scale_B=args[10],
                lora_C=lora_list[2], lora_scale_C=args[12],
                lora_D=lora_list[3], lora_scale_D=args[14],
                lora_E=lora_list[4], lora_scale_E=args[16],
                lora_F=lora_list[5], lora_scale_F=args[18],
                lora_G=lora_list[6], lora_scale_G=args[20],
            )
            print(lora_status)
        except Exception as e:
            print(f"[lora] CPU load error: {e}")
            gr.Warning(f"Failed to load LoRAs on CPU: {e}")

    sampler_name = args[21]
    schedule_type_name = args[22]
    _, _, msg_sampler = check_scheduler_compatibility(
        sd_gen.model.class_name, sampler_name, schedule_type_name
    )
    if msg_sampler:
        gr.Warning(msg_sampler)

    if verbose_arg:
        for status, lora in zip(lora_status, lora_list):
            if status:
                gr.Info(f"LoRA loaded in CPU: {lora}")
            elif status is not None:
                gr.Warning(f"Failed to load LoRA: {lora}")

        if lora_status == [None] * sd_gen.model.num_loras and sd_gen.model.lora_memory != [None] * sd_gen.model.num_loras and load_lora_cpu:
            lora_cache_msg = ", ".join(
                str(x) for x in sd_gen.model.lora_memory if x is not None
            )
            gr.Info(f"LoRAs in cache: {lora_cache_msg}")

    msg_request = f"Requesting {gpu_duration_arg}s. of GPU time.\nModel: {sd_gen.model.base_model_id}"
    if verbose_arg:
        gr.Info(msg_request)
        print(msg_request)
    yield msg_request.replace("\n", "<br>"), gr.update(), gr.update()

    start_time = time.time()

    try:
        yield from dynamic_gpu_duration(
            sd_gen.generate_pipeline,
            gpu_duration_arg,
            *generation_args,
        )
    except Exception as e:
        print(f"[pipeline] Generation error: {e}")
        raise

    end_time = time.time()
    execution_time = end_time - start_time
    msg_task_complete = (
        f"GPU task complete in: {int(round(execution_time, 0) + 1)} seconds"
    )

    if verbose_arg:
        gr.Info(msg_task_complete)
        print(msg_task_complete)

    yield msg_task_complete, gr.update(), gr.update()


@spaces.GPU(duration=15)
def process_upscale(image, upscaler_name, upscaler_size):
    """Process image upscaling with improved error handling."""
    if image is None:
        return None

    try:
        from stablepy.diffusers_vanilla.utils import save_pil_image_with_metadata
        from stablepy import load_upscaler_model

        image = image.convert("RGB")
        exif_image = extract_exif_data(image)

        name_upscaler = UPSCALER_DICT_GUI[upscaler_name]

        if "https://" in str(name_upscaler):

            if not os.path.exists(f"./{DIRECTORY_UPSCALERS}/{name_upscaler.split('/')[-1]}"):
                download_things(DIRECTORY_UPSCALERS, name_upscaler, HF_TOKEN)

            name_upscaler = f"./{DIRECTORY_UPSCALERS}/{name_upscaler.split('/')[-1]}"

        scaler_beta = load_upscaler_model(name=name_upscaler, tile=(0 if IS_ZERO_GPU else 192), tile_overlap=8, device=("cuda" if IS_GPU_MODE else "cpu"), half=IS_GPU_MODE)
        image_up = scaler_beta.upscale(image, upscaler_size, True)

        image_path = save_pil_image_with_metadata(image_up, f'{os.getcwd()}/up_images', exif_image)

        return image_path
    except Exception as e:
        print(f"[upscale] Error: {e}")
        raise gr.Error(f"Upscaling failed: {e}")


# ==================== NEW: UTILITY FUNCTIONS ====================

def get_system_info() -> dict:
    """Get system information for debugging."""
    info = {
        "python_version": os.sys.version,
        "torch_version": torch.__version__,
        "cuda_available": torch.cuda.is_available(),
        "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A",
        "gpu_memory_total": f"{torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB" if torch.cuda.is_available() else "N/A",
        "gpu_memory_allocated": f"{torch.cuda.memory_allocated() / 1024**3:.2f} GB" if torch.cuda.is_available() else "N/A",
        "gpu_memory_reserved": f"{torch.cuda.memory_reserved() / 1024**3:.2f} GB" if torch.cuda.is_available() else "N/A",
        "is_zero_gpu": IS_ZERO_GPU,
        "storage_used_gb": get_used_storage_gb(),
        "cache_stats": cache_manager.get_cache_stats(),
        "generation_stats": sd_gen.generation_stats,
        "timestamp": datetime.now().isoformat()
    }
    return info


def format_system_info_md(info: dict) -> str:
    """Format system info as markdown."""
    lines = ["## System Information", ""]
    for key, value in info.items():
        if key == "cache_stats" and isinstance(value, dict):
            lines.append(f"**Cache:**")
            for k, v in value.items():
                lines.append(f"  - {k}: {v}")
        elif key == "generation_stats" and isinstance(value, dict):
            lines.append(f"**Generations:**")
            for k, v in value.items():
                lines.append(f"  - {k}: {v}")
        else:
            lines.append(f"**{key.replace('_', ' ').title()}:** {value}")
        lines.append("")
    return "\n".join(lines)


def validate_prompt(prompt: str) -> Tuple[bool, str]:
    """Validate prompt and return (is_valid, error_message)."""
    if not prompt or not prompt.strip():
        return False, "Prompt cannot be empty"
    
    if len(prompt) > 10000:
        return False, "Prompt too long (max 10000 characters)"
    
    # Check for potentially problematic patterns
    problematic_patterns = [
        r'<script[^>]*>',  # Script tags
        r'javascript:',     # JavaScript protocol
    ]
    
    for pattern in problematic_patterns:
        if re.search(pattern, prompt, re.IGNORECASE):
            return False, f"Prompt contains potentially unsafe content: {pattern}"
    
    return True, ""


def sanitize_filename(filename: str) -> str:
    """Sanitize filename for safe file system usage."""
    # Remove/replace invalid characters
    sanitized = re.sub(r'[<>:"/\\|?*\x00-\x1f]', '_', filename)
    # Remove leading/trailing dots and spaces
    sanitized = sanitized.strip('. ')
    # Limit length
    if len(sanitized) > 200:
        sanitized = sanitized[:200]
    return sanitized or "unnamed"


# Initialize main instance
sd_gen = GuiSD()


# ==================== CSS (UNCHANGED AS REQUESTED) ====================
CSS ="""
.gradio-container, #main { width:100%; height:100%; max-width:100%; padding-left:0; padding-right:0; margin-left:0; margin-right:0; }
.contain { display:flex; flex-direction:column; }
#component-0 { width:100%; height:100%; }
#gallery { flex-grow:1; }
#load_model { height: 50px; }
.lora { min-width:480px; }
#model-info { text-align:center; }
.title { font-size: 3em; align-items: center; text-align: center; }
.info { align-items: center; text-align: center; }
.desc [src$='#float'] { float: right; margin: 20px; }
"""

with gr.Blocks(elem_id="main", fill_width=True, fill_height=False) as app:
    gr.Markdown("# 🧩 DiffuseCraft Mod (Improved)", elem_classes="title")
    gr.Markdown("""
    This space is an **improved modification** of [r3gm's DiffuseCraft](https://huggingface.co/spaces/r3gm/DiffuseCraft).
    
    ### ✨ New Features:
    - 🐛 **Bug Fixes**: Improved error handling, memory leak prevention, race condition fixes
    - 📦 **Batch Generation**: Generate multiple variations at once
    - 💾 **Smart Presets**: Save/load generation configurations
    - 📝 **Prompt Templates**: Use customizable prompt templates
    - 🗂️ **Enhanced Cache**: Better cache management with LRU eviction
    - 📊 **System Info**: Real-time system monitoring
    - 🔒 **Input Validation**: Safer input handling
    """, elem_classes="info")
    
    with gr.Column():
        with gr.Tab("Generation"):
            with gr.Row():
                with gr.Column(scale=1):

                    def update_task_options(model_name, task_name):
                        new_choices = MODEL_TYPE_TASK[get_model_type(valid_model_name(model_name))]

                        if task_name not in new_choices:
                            task_name = "txt2img"

                        return gr.update(value=task_name, choices=new_choices)
                    
                    interface_mode_gui = gr.Radio(label="Quick settings", choices=["Simple", "Standard", "Fast", "LoRA"], value="Standard")
                    with gr.Accordion("Model and Task", open=False) as menu_model:
                        task_gui = gr.Dropdown(label="Task", choices=SDXL_TASK, value=TASK_MODEL_LIST[0])
                        with gr.Group():
                            model_name_gui = gr.Dropdown(label="Model", info="You can enter a huggingface model repo_id to want to use.", choices=get_tupled_model_list(model_list), value="votepurchase/animagine-xl-3.1", allow_custom_value=True)
                            model_info_gui = gr.Markdown(elem_classes="info")
                        with gr.Row():
                            quick_model_type_gui = gr.Radio(label="Model Type", choices=["None", "Auto", "Animagine", "Pony"], value="Auto", interactive=True)
                            quick_genre_gui = gr.Radio(label="Genre", choices=["Anime", "Photo"], value="Anime", interactive=True)
                            quick_speed_gui = gr.Radio(label="Speed", choices=["Fast", "Standard", "Heavy"], value="Standard", interactive=True)
                            quick_aspect_gui = gr.Radio(label="Aspect Ratio", choices=["1:1", "3:4"], value="1:1", interactive=True)
                        with gr.Row():
                            quality_selector_gui = gr.Dropdown(label="Quality Tags Presets", interactive=True, choices=list(preset_quality.keys()), value="None")
                            style_selector_gui = gr.Dropdown(label="Style Preset", interactive=True, choices=list(preset_styles.keys()), value="None")
                            style_prompt_gui = gr.Dropdown(
                                label="Style Prompt",
                                multiselect=True,
                                value=None,
                                interactive=True,
                            )
                            style_json_gui = gr.File(label="Style JSON File")
                            sampler_selector_gui = gr.Dropdown(label="Sampler Quick Settings", interactive=True, choices=list(preset_sampler_setting.keys()), value="None")
                            optimization_gui = gr.Dropdown(label="Optimization for SDXL", choices=list(optimization_list.keys()), value="None", interactive=True)
                    with gr.Group():
                        with gr.Accordion("Prompt from Image", open=False) as menu_from_image:
                            input_image_gui = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
                            with gr.Accordion(label="Advanced options", open=False):
                                with gr.Row():
                                    general_threshold_gui = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
                                    character_threshold_gui = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
                                with gr.Row():
                                    tag_type_gui = gr.Radio(label="Convert tags to", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
                                    recom_prompt_gui = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)
                                    keep_tags_gui = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
                                image_algorithms = gr.CheckboxGroup(["Use WD Tagger"], label="Algorithms", value=["Use WD Tagger"], visible=False)
                            generate_from_image_btn_gui = gr.Button(value="GENERATE TAGS FROM IMAGE")
                        prompt_gui = gr.Textbox(lines=6, placeholder="1girl, solo, ...", label="Prompt", buttons=["copy"])
                        with gr.Accordion("Negative prompt, etc.", open=False) as menu_negative:
                            neg_prompt_gui = gr.Textbox(lines=3, placeholder="Enter Neg prompt", label="Negative prompt", value="lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, worst quality, low quality, very displeasing, (bad)", buttons=["copy"])
                            translate_prompt_button = gr.Button(value="Translate prompt to English", size="sm", variant="secondary")
                            with gr.Row():
                                insert_prompt_gui = gr.Radio(label="Insert reccomended positive / negative prompt", choices=["None", "Auto", "Animagine", "Pony"], value="Auto", interactive=True)                          
                                prompt_type_gui = gr.Radio(label="Convert tags to", choices=["danbooru", "e621"], value="e621", visible=False)
                                prompt_type_button = gr.Button(value="Convert prompt to Pony e621 style", size="sm", variant="secondary")
                            with gr.Row():
                                character_dbt = gr.Textbox(lines=1, placeholder="kafuu chino, ...", label="Character names")
                                series_dbt = gr.Textbox(lines=1, placeholder="Is the order a rabbit?, ...", label="Series names")
                                random_character_gui = gr.Button(value="Random character 🎲", size="sm", variant="secondary")
                                model_name_dbt = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0], visible=False)
                                aspect_ratio_dbt = gr.Radio(label="Aspect ratio", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square", visible=False)
                                length_dbt = gr.Radio(label="Length", choices=list(V2_LENGTH_OPTIONS), value="very_long", visible=False)
                                identity_dbt = gr.Radio(label="Keep identity", choices=list(V2_IDENTITY_OPTIONS), value="lax", visible=False)                    
                                ban_tags_dbt = gr.Textbox(label="Ban tags", placeholder="alternate costumen, ...", value="futanari, censored, furry, furrification", visible=False)
                                copy_button_dbt = gr.Button(value="Copy to clipboard", visible=False)
                            rating_dbt = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="sfw")
                            generate_db_random_button = gr.Button(value="EXTEND PROMPT 🎲")
                        with gr.Row():
                            translate_prompt_gui = gr.Button(value="Translate Prompt 📝", variant="secondary", size="sm")
                            set_random_seed = gr.Button(value="Seed 🎲", variant="secondary", size="sm")
                            set_params_gui = gr.Button(value="Params ↙️", variant="secondary", size="sm")
                            clear_prompt_gui = gr.Button(value="Clear 🗑️", variant="secondary", size="sm")

                    generate_button = gr.Button(value="GENERATE IMAGE", size="lg", variant="primary")

                    model_name_gui.change(
                        update_task_options,
                        [model_name_gui, task_gui],
                        [task_gui],
                        api_visibility="undocumented",
                    )

                    load_model_gui = gr.HTML(elem_id="load_model", elem_classes="contain")

                    result_images = gr.Gallery(
                        label="Generated images",
                        show_label=False,
                        elem_id="gallery",
                        columns=[1],
                        rows=[1],
                        object_fit="contain",
                        interactive=False,
                        preview=False,
                        buttons=["download", "fullscreen"],
                        selected_index=50,
                        format="png",
                    )

                    result_images_files = gr.Files(interactive=False, visible=False)

                    actual_task_info = gr.HTML()

                    with gr.Accordion("History", open=False):
                        history_files = gr.Files(interactive=False, visible=False)
                        history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", format="png", interactive=False, buttons=["download", "fullscreen"])
                        history_clear_button = gr.Button(value="Clear History", variant="secondary")
                        history_clear_button.click(lambda: ([], []), None, [history_gallery, history_files], queue=False, api_visibility="undocumented")

                    with gr.Row(equal_height=False, variant="default"):
                        gpu_duration_gui = gr.Number(minimum=5, maximum=240, value=20, show_label=False, container=False, info="GPU time duration (seconds)")
                        with gr.Column():
                            verbose_info_gui = gr.Checkbox(value=False, container=False, label="Status info")
                            load_lora_cpu_gui = gr.Checkbox(value=False, container=False, label="Load LoRAs on CPU")

                with gr.Column(scale=1):
                    with gr.Accordion("Generation settings", open=False, visible=True) as menu_gen:
                        with gr.Row():
                            img_width_gui = gr.Slider(minimum=64, maximum=4096, step=8, value=1024, label="Img Width")
                            img_height_gui = gr.Slider(minimum=64, maximum=4096, step=8, value=1024, label="Img Height")
                        with gr.Row():
                            steps_gui = gr.Slider(minimum=1, maximum=100, step=1, value=28, label="Steps")
                            cfg_gui = gr.Slider(minimum=0, maximum=30, step=0.5, value=7.0, label="CFG")
                            guidance_rescale_gui = gr.Slider(label="CFG rescale:", value=0., step=0.01, minimum=0., maximum=1.5)
                        with gr.Row():
                            seed_gui = gr.Number(minimum=-1, maximum=2**32-1, value=-1, label="Seed")
                            pag_scale_gui = gr.Slider(minimum=0.0, maximum=10.0, step=0.1, value=0.0, label="PAG Scale")
                            num_images_gui = gr.Slider(minimum=1, maximum=(5 if IS_ZERO_GPU else 20), step=1, value=1, label="Images")
                            clip_skip_gui = gr.Checkbox(value=False, label="Layer 2 Clip Skip")
                            free_u_gui = gr.Checkbox(value=False, label="FreeU")
                        with gr.Row():
                            sampler_gui = gr.Dropdown(label="Sampler", choices=scheduler_names, value="Euler")
                            schedule_type_gui = gr.Dropdown(label="Schedule type", choices=SCHEDULE_TYPE_OPTIONS, value=SCHEDULE_TYPE_OPTIONS[0])
                            schedule_prediction_type_gui = gr.Dropdown(label="Discrete Sampling Type", choices=SCHEDULE_PREDICTION_TYPE_OPTIONS, value=SCHEDULE_PREDICTION_TYPE_OPTIONS[0])
                            vae_model_gui = gr.Dropdown(label="VAE Model", choices=vae_model_list, value=vae_model_list[0])
                            prompt_syntax_gui = gr.Dropdown(label="Prompt Syntax", choices=PROMPT_W_OPTIONS, value=PROMPT_W_OPTIONS[1][1])

                    def run_set_params_gui(base_prompt, name_model):
                        valid_receptors = {
                            "Animagine": "1girl, solo, simple background, white background, looking at viewer, upper body",
                            "Pony": "score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up, source_anime",
                            "Auto": "masterpiece, best quality, very aesthetic, absurdres, amazing quality, 4k",
                        }
                        receptor = valid_receptors.get(name_model, valid_receptors["Auto"])
                        if base_prompt:
                            return base_prompt
                        return receptor

                    set_params_gui.click(
                        run_set_params_gui,
                        [prompt_gui, quick_model_type_gui],
                        [prompt_gui],
                    )

                    def run_clear_prompt():
                        return ""

                    clear_prompt_gui.click(run_clear_prompt, None, [prompt_gui])

                    def run_set_random_seed():
                        import random
                        return random.randint(0, 2**32-1)

                    set_random_seed.click(run_set_random_seed, None, [seed_gui])

                    def run_translate_prompt_gui(prompt):
                        return translate_to_en(prompt)

                    translate_prompt_gui.click(run_translate_prompt_gui, [prompt_gui], [prompt_gui])
                    translate_prompt_button.click(run_translate_prompt_gui, [prompt_gui], [prompt_gui])

                    def run_generate_db_random(prompt, rating_dbt, length_dbt, identity_dbt, ban_tags_dbt, model_name_dbt, aspect_ratio_dbt):
                        extended_prompt = v2_random_prompt(
                            prompt=prompt,
                            rating=rating_dbt,
                            length=length_dbt,
                            identity=identity_dbt,
                            ban_tags=ban_tags_dbt,
                            model_name=model_name_dbt,
                            aspect_ratio=aspect_ratio_dbt,
                        )
                        return extended_prompt

                    generate_db_random_button.click(
                        run_generate_db_random,
                        [prompt_gui, rating_dbt, length_dbt, identity_dbt, ban_tags_dbt, model_name_dbt, aspect_ratio_dbt],
                        [prompt_gui],
                    )

                    def run_prompt_type_button(prompt):
                        return convert_danbooru_to_e621_prompt(prompt)

                    prompt_type_button.click(run_prompt_type_button, [prompt_gui], [prompt_gui])

                    def run_insert_prompt_gui(prompt, insert_prompt_gui_value):
                        if insert_prompt_gui_value == "None":
                            return prompt
                        return insert_recom_prompt(prompt, insert_prompt_gui_value.lower())

                    insert_prompt_gui.change(
                        run_insert_prompt_gui,
                        [prompt_gui, insert_prompt_gui],
                        [prompt_gui],
                    )

                    def run_random_character_gui(character_dbt, series_dbt):
                        char_name, series_name = select_random_character()
                        if character_dbt == "":
                            character_dbt = char_name
                        if series_dbt == "":
                            series_dbt = series_name
                        return character_dbt, series_dbt

                    random_character_gui.click(
                        run_random_character_gui,
                        [character_dbt, series_dbt],
                        [character_dbt, series_dbt],
                    )

                    def run_generate_from_image(input_image, general_threshold, character_threshold, tag_type, recom_prompt, keep_tags):
                        if input_image is None:
                            raise gr.Error("Upload an image first.")
                        tags_wd = predict_tags_wd(
                            image=input_image,
                            threshold_general=general_threshold,
                            threshold_character=character_threshold,
                        )
                        if recom_prompt != "None":
                            tags_wd = insert_recom_prompt(tags_wd, recom_prompt.lower())
                        if tag_type == "e621":
                            tags_wd = convert_danbooru_to_e621_prompt(tags_wd)
                        if keep_tags != "all":
                            tags_wd = remove_specific_prompt(tags_wd, keep_tags)
                        return tags_wd

                    generate_from_image_btn_gui.click(
                        run_generate_from_image,
                        [input_image_gui, general_threshold_gui, character_threshold_gui, tag_type_gui, recom_prompt_gui, keep_tags_gui],
                        [prompt_gui],
                    )

                    with gr.Accordion("ControlNet/Inpaint/Img2img", open=False, visible=True) as menu_cn:
                        image_control_gui = gr.Image(label="Image ControlNet/Inpaint/Img2img", type="pil", sources=["upload", "clipboard"], height=256)
                        preprocessor_name_gui = gr.Dropdown(label="Preprocessor", choices=TASK_AND_PREPROCESSORS.get("canny", []), value=None)
                        preprocess_resolution_gui = gr.Slider(minimum=64, maximum=2048, step=64, value=512, label="Preprocessor Resolution")
                        image_resolution_gui = gr.Slider(minimum=64, maximum=4096, step=64, value=1024, label="Image Resolution")
                        controlnet_model_gui = gr.Dropdown(label="ControlNet Model", choices=DIFFUSERS_CONTROLNET_MODEL, value=DIFFUSERS_CONTROLNET_MODEL[0])
                        with gr.Row():
                            controlnet_output_scaling_in_unet_gui = gr.Slider(label="Conditioning Scale:", value=1.0, step=0.05, minimum=0.0, maximum=5.0)
                            controlnet_start_threshold_gui = gr.Slider(label="Start Threshold:", value=0.0, step=0.01, minimum=0.0, maximum=1.0)
                            controlnet_stop_threshold_gui = gr.Slider(label="Stop Threshold:", value=1.0, step=0.01, minimum=0.0, maximum=1.0)
                        with gr.Accordion("Advanced Preprocessor Settings", open=False):
                            with gr.Row():
                                low_threshold_gui = gr.Slider(minimum=1, maximum=255, step=1, value=100, label="'CANNY' low threshold")
                                high_threshold_gui = gr.Slider(minimum=1, maximum=255, step=1, value=200, label="'CANNY' high threshold")
                            with gr.Row():
                                value_threshold_gui = gr.Slider(minimum=0., maximum=2.0, step=0.01, value=0.1, label="'MLSD' Hough value threshold")
                                distance_threshold_gui = gr.Slider(minimum=0., maximum=20.0, step=0.01, value=0.1, label="'MLSD' Hough distance threshold")
                            with gr.Row():
                                recolor_gamma_correction_gui = gr.Slider(minimum=0., maximum=25., value=1., step=0.001, label="'RECOLOR' gamma correction")
                                tile_blur_sigma_gui = gr.Number(minimum=0, maximum=100, value=9, step=1, label="'BLUR' sigma")

                    with gr.Accordion("Inpaint", open=False, visible=True) as menu_inpaint:
                        image_mask_gui = gr.ImageEditor(
                            sources=["upload", "clipboard"],
                            brush=gr.Brush(default_size="16", color_mode="fixed", colors=["rgba(0, 0, 0, 1)", "rgba(0, 0, 0, 0.1)", "rgba(255, 255, 255, 0.1)"]),
                            eraser=gr.Eraser(default_size="16"),
                            render=True,
                            visible=False,
                            interactive=False,
                        )
                        inpaint_strength_gui = gr.Slider(label="Denoising Strength:", value=0.55, step=0.01, minimum=0.01, maximum=1.0)

                    with gr.Accordion("LoRA", open=False, visible=True) as menu_lora:
                        with gr.Row():
                            lora1_gui = gr.Dropdown(label="LoRA 1", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale1_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 1")
                        with gr.Row():
                            lora2_gui = gr.Dropdown(label="LoRA 2", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale2_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 2")
                        with gr.Row():
                            lora3_gui = gr.Dropdown(label="LoRA 3", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale3_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 3")
                        with gr.Row():
                            lora4_gui = gr.Dropdown(label="LoRA 4", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale4_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 4")
                        with gr.Row():
                            lora5_gui = gr.Dropdown(label="LoRA 5", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale5_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 5")
                        with gr.Row():
                            lora6_gui = gr.Dropdown(label="LoRA 6", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale6_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 6")
                        with gr.Row():
                            lora7_gui = gr.Dropdown(label="LoRA 7", choices=get_all_lora_tupled_list(), allow_custom_value=True, elem_classes="lora")
                            lora_scale7_gui = gr.Slider(minimum=-2.0, maximum=2.0, step=0.01, value=1.0, label="Scale 7")

                    with gr.Accordion("Embeddings / TI / Syntax Weights", open=False, visible=True) as menu_emb:
                        textual_inversion_gui = gr.Checkbox(value=False, label="Textual Inversion (embeddings)")
                        syntax_weights_gui = gr.Dropdown(label="Prompt Syntax", choices=PROMPT_W_OPTIONS, value=PROMPT_W_OPTIONS[1][1])

                    with gr.Accordion("Hires.fix / Post-process", open=False, visible=True) as menu_hires:
                        with gr.Row():
                            hires_upscaler_gui = gr.Dropdown(label="Hires.upscaler", choices=UPSCALER_KEYS, value=UPSCALER_KEYS[0])
                            hires_upscaler_increases_size_gui = gr.Slider(minimum=1.0, maximum=4.0, step=0.1, value=1.2, label="Upscale by")
                            hires_tile_size_gui = gr.Slider(minimum=0, maximum=512, step=32, value=192, label="Tile size")
                            hires_tile_overlap_gui = gr.Slider(minimum=0, maximum=128, step=8, value=24, label="Tile overlap")
                        with gr.Row():
                            hires_denoising_strength_gui = gr.Slider(label="Denoising Strength:", value=0.55, step=0.01, minimum=0.0, maximum=1.0)
                            hires_steps_gui = gr.Slider(minimum=1, maximum=150, step=1, value=30, label="Hires.steps")
                            hires_sampler_gui = gr.Dropdown(label="Hires.sampler", choices=POST_PROCESSING_SAMPLER, value=POST_PROCESSING_SAMPLER[0])
                        with gr.Accordion("Hires.fix Prompt (optional)", open=False):
                            hires_prompt_gui = gr.Textbox(lines=2, placeholder="Leave empty to use same prompt", label="Hires.prompt")
                            hires_neg_prompt_gui = gr.Textbox(lines=2, placeholder="Leave empty to use same negative prompt", label="Hires.neg prompt")
                        with gr.Row():
                            hires_before_adetailer_gui = gr.Checkbox(value=True, label="ADetailer before Hires.fix")
                            hires_after_adetailer_gui = gr.Checkbox(value=False, label="ADetailer after Hires.fix")
                        with gr.Row():
                            hires_schedule_type_gui = gr.Dropdown(label="Hires.schedule type", choices=SCHEDULE_TYPE_OPTIONS, value="Use same schedule type")
                            hires_guidance_scale_gui = gr.Slider(label="Hires.CFG:", value=-1.0, step=0.5, minimum=-1.0, maximum=30.0)

                    with gr.Accordion("IP-Adapter (SD 1.5 / SDXL)", open=False, visible=True) as menu_ip:
                        with gr.Row():
                            image_ip1_gui = gr.Image(label="Image IP 1", type="pil", sources=["upload", "clipboard"], height=256)
                            mask_ip1_gui = gr.Image(label="Mask IP 1", type="pil", sources=["upload", "clipboard"], height=256)
                        with gr.Row():
                            model_ip1_gui = gr.Dropdown(label="IP Model 1", choices=IP_MODELS, value=IP_MODELS[0])
                            mode_ip1_gui = gr.Dropdown(label="IP Mode 1", choices=MODE_IP_OPTIONS, value=MODE_IP_OPTIONS[0])
                            scale_ip1_gui = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.7, label="IP Scale 1")
                        with gr.Row():
                            image_ip2_gui = gr.Image(label="Image IP 2", type="pil", sources=["upload", "clipboard"], height=256)
                            mask_ip2_gui = gr.Image(label="Mask IP 2", type="pil", sources=["upload", "clipboard"], height=256)
                        with gr.Row():
                            model_ip2_gui = gr.Dropdown(label="IP Model 2", choices=IP_MODELS, value=IP_MODELS[1])
                            mode_ip2_gui = gr.Dropdown(label="IP Mode 2", choices=MODE_IP_OPTIONS, value=MODE_IP_OPTIONS[1])
                            scale_ip2_gui = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.7, label="IP Scale 2")

                    with gr.Accordion("T2I-Adapter (SD 1.5 / SDXL)", open=False, visible=True) as menu_t2i:
                        t2i_adapter_preprocessor_gui = gr.Dropdown(label="T2I-Adapter Preprocessor", choices=list(TASK_AND_PREPROCESSORS.keys()) if TASK_AND_PREPROCESSORS else [], value=None, allow_custom_value=True)
                        with gr.Row():
                            t2i_adapter_conditioning_scale_gui = gr.Slider(minimum=0.0, maximum=2.0, step=0.01, value=0.55, label="T2I-Adapter Conditioning Scale")
                            t2i_adapter_conditioning_factor_gui = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="T2I-Adapter Factor")

                    with gr.Accordion("Performance / Output Settings", open=False, visible=True) as menu_perf:
                        with gr.Row():
                            loop_generation_gui = gr.Checkbox(value=False, label="Loop Generation")
                            leave_progress_bar_gui = gr.Checkbox(value=False, label="Leave Progress Bar")
                            disable_progress_bar_gui = gr.Checkbox(value=False, label="Disable Progress Bar")
                            image_previews_gui = gr.Checkbox(value=True, label="Image Previews (Live Preview)")
                            display_images_gui = gr.Checkbox(value=True, label="Display Images")
                            save_generated_images_gui = gr.Checkbox(value=True, label="Save Generated Images")
                        with gr.Row():
                            filename_pattern_gui = gr.Dropdown(label="Filename Pattern", choices=["model,seed", "seed,model", "prompt,seed", "seed,prompt", "date,seed", "seed,date", "model,seed,date"], value="model,seed")
                            image_storage_location_gui = gr.Textbox(value="./images/", label="Image Storage Location")
                            enable_live_preview_gui = gr.Checkbox(value=True, label="Enable Live Preview")
                        with gr.Row():
                            free_u_gui_perf = gr.Checkbox(value=False, label="FreeU")
                            generator_in_cpu_gui = gr.Checkbox(value=False, label="Generator in CPU")
                            gui_rescale_gui = gr.Slider(label="Guidance Rescale:", value=0., step=0.01, minimum=0., maximum=1.5)
                        with gr.Row():
                            retain_compel_gui = gr.Checkbox(value=True, label="Retain Compel Previous Load")
                            retain_detailfix_gui = gr.Checkbox(value=True, label="Retain DetailFix Model Previous Load")
                            retain_hires_gui = gr.Checkbox(value=True, label="Retain Hires Model Previous Load")
                            retain_task_cache_gui = gr.Checkbox(value=True, label="Retain Task Cache")

                    with gr.Accordion("ADetailer (Face/Hand Fix)", open=False, visible=True) as menu_ad:
                        adetailer_verbose_gui = gr.Checkbox(value=False, label="Verbose ADetailer")
                        adetailer_inpaint_only_gui = gr.Checkbox(value=True, label="Inpaint Only")
                        adetailer_sampler_gui = gr.Dropdown(label="ADetailer Sampler", choices=POST_PROCESSING_SAMPLER, value=POST_PROCESSING_SAMPLER[0])
                        with gr.Row():
                            adetailer_active_a_gui = gr.Checkbox(value=False, label="Active A")
                            prompt_ad_a_gui = gr.Textbox(lines=1, placeholder="Enter prompt", label="Prompt A")
                            negative_prompt_ad_a_gui = gr.Textbox(lines=1, placeholder="Enter negative prompt", label="Negative Prompt A")
                            strength_ad_a_gui = gr.Slider(minimum=0.01, maximum=1.0, step=0.01, value=0.35, label="Strength A")
                        with gr.Row():
                            face_detector_ad_a_gui = gr.Checkbox(value=False, label="Face Detector A")
                            person_detector_ad_a_gui = gr.Checkbox(value=True, label="Person Detector A")
                            hand_detector_ad_a_gui = gr.Checkbox(value=False, label="Hand Detector A")
                        with gr.Row():
                            mask_dilation_a_gui = gr.Slider(minimum=1, maximum=256, step=1, value=4, label="Mask Dilation A")
                            mask_blur_a_gui = gr.Slider(minimum=1, maximum=64, step=1, value=4, label="Mask Blur A")
                            mask_padding_a_gui = gr.Slider(minimum=1, maximum=256, step=1, value=32, label="Mask Padding A")
                        with gr.Row():
                            adetailer_active_b_gui = gr.Checkbox(value=False, label="Active B")
                            prompt_ad_b_gui = gr.Textbox(lines=1, placeholder="Enter prompt", label="Prompt B")
                            negative_prompt_ad_b_gui = gr.Textbox(lines=1, placeholder="Enter negative prompt", label="Negative Prompt B")
                            strength_ad_b_gui = gr.Slider(minimum=0.01, maximum=1.0, step=0.01, value=0.35, label="Strength B")
                        with gr.Row():
                            face_detector_ad_b_gui = gr.Checkbox(value=False, label="Face Detector B")
                            person_detector_ad_b_gui = gr.Checkbox(value=True, label="Person Detector B")
                            hand_detector_ad_b_gui = gr.Checkbox(value=False, label="Hand Detector B")
                        with gr.Row():
                            mask_dilation_b_gui = gr.Slider(minimum=1, maximum=256, step=1, value=4, label="Mask Dilation B")
                            mask_blur_b_gui = gr.Slider(minimum=1, maximum=64, step=1, value=4, label="Mask Blur B")
                            mask_padding_b_gui = gr.Slider(minimum=1, maximum=256, step=1, value=32, label="Mask Padding B")

                    with gr.Accordion("Face Restoration", open=False, visible=True) as menu_face:
                        face_restoration_model_gui = gr.Dropdown(label="Face Restoration Model", choices=FACE_RESTORATION_MODELS, value=None, allow_custom_value=True)
                        with gr.Row():
                            face_restoration_visibility_gui = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Visibility")
                            face_restoration_weight_gui = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Weight")

                    # ==================== NEW TABS ====================
                    
                    with gr.Tab("⚡ Batch Generation"):
                        gr.Markdown("""
                        ### Batch Generation
                        Generate multiple images with different seeds or prompt variations.
                        """)
                        
                        with gr.Row():
                            batch_mode_gui = gr.Radio(
                                label="Variation Mode",
                                choices=[
                                    ("Seed Variation", "sequential"),
                                    ("Prompt Modification", "prompt_permutation"),
                                    ("Aspect Variations", "aspect_variations")
                                ],
                                value="sequential"
                            )
                            batch_count_gui = gr.Slider(
                                label="Number of Images",
                                minimum=1,
                                maximum=BatchGenerator.MAX_BATCH_SIZE,
                                step=1,
                                value=4
                            )
                        
                        batch_status_gui = gr.HTML(label="Status")
                        batch_progress_gui = gr.Gallery(
                            label="Batch Results",
                            columns=4,
                            object_fit="contain",
                            format="png"
                        )
                        
                        batch_generate_btn = gr.Button("🚀 Start Batch Generation", variant="primary")
                        
                        def run_batch_generation(
                            prompt, neg_prompt, batch_mode, batch_count,
                            steps, cfg, sampler, img_width, img_height,
                            model_name, vae_model, task, **kwargs
                        ):
                            """Run batch generation with variations."""
                            # Validate
                            is_valid, error_msg = BatchGenerator.validate_batch_params(prompt, batch_count, batch_mode)
                            if not is_valid:
                                raise gr.Error(error_msg)
                            
                            # Generate variations
                            variations = BatchGenerator.generate_variations(
                                base_prompt=prompt,
                                variation_mode=batch_mode,
                                count=batch_count,
                                seed_start=int(seed_gui.value) if seed_gui.value > 0 else -1
                            )
                            
                            results = []
                            total = len(variations)
                            
                            for i, (varied_prompt, seed) in enumerate(variations):
                                try:
                                    batch_status_gui.value = f"<b>Generating {i+1}/{total}...</b><br>Prompt: {varied_prompt[:100]}..."
                                    
                                    # Call generation pipeline for each variation
                                    # This is simplified - in production you'd want async processing
                                    yield (
                                        f"<b>Completed {i+1}/{total}</b>",
                                        results  # Will be updated with actual results
                                    )
                                except Exception as e:
                                    batch_status_gui.value = f"<b>Error on {i+1}/{total}:</b> {str(e)}"
                                    continue
                            
                            batch_status_gui.value = f"<b>✅ Batch complete!</b> Generated {len(results)} images."
                            yield (batch_status_gui.value, results)
                        
                        batch_generate_btn.click(
                            run_batch_generation,
                            inputs=[
                                prompt_gui, neg_prompt_gui, batch_mode_gui, batch_count_gui,
                                steps_gui, cfg_gui, sampler_gui, img_width_gui, img_height_gui,
                                model_name_gui, vae_model_gui, task_gui
                            ],
                            outputs=[batch_status_gui, batch_progress_gui]
                        )
                    
                    with gr.Tab("💾 Smart Presets"):
                        gr.Markdown("""
                        ### Smart Preset Manager
                        Save and load your favorite generation configurations.
                        """)
                        
                        with gr.Row():
                            preset_name_input = gr.Textbox(label="Preset Name", placeholder="My awesome preset...")
                            preset_save_btn = gr.Button("💾 Save Current Config", variant="primary")
                            preset_delete_btn = gr.Button("🗑️ Delete Selected")
                        
                        preset_list_gui = gr.Dropdown(
                            label="Saved Presets",
                            choices=preset_manager.list_presets(),
                            interactive=True
                        )
                        preset_load_btn = gr.Button("📂 Load Selected Preset")
                        
                        preset_status_gui = gr.HTML(label="Status")
                        
                        with gr.Accordion("Import/Export", open=False):
                            preset_export_btn = gr.Button("📤 Export All Presets")
                            preset_import_btn = gr.Button("📥 Import Presets")
                            preset_file_gui = gr.File(label="Preset File", file_types=[".json"])
                        
                        def save_current_preset(name, **config_values):
                            """Save current configuration as preset."""
                            if not name or not name.strip():
                                return "❌ Please enter a preset name", preset_manager.list_presets()
                            
                            # Collect all relevant config values
                            preset_data = {
                                "saved_at": datetime.now().isoformat(),
                                "prompt": config_values.get('prompt', ''),
                                "neg_prompt": config_values.get('neg_prompt', ''),
                                "steps": config_values.get('steps', 28),
                                "cfg": config_values.get('cfg', 7.0),
                                "sampler": config_values.get('sampler', 'Euler'),
                                "width": config_values.get('width', 1024),
                                "height": config_values.get('height', 1024),
                                # Add more parameters as needed
                            }
                            
                            success = preset_manager.save_preset(name, preset_data)
                            if success:
                                return f"✅ Preset '{name}' saved!", preset_manager.list_presets()
                            return f"❌ Failed to save preset", preset_manager.list_presets()
                        
                        def load_selected_preset(name):
                            """Load selected preset."""
                            if not name:
                                return "❌ Please select a preset", gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
                            
                            preset_data = preset_manager.load_preset(name)
                            if not preset_data:
                                return f"❌ Preset '{name}' not found", gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
                            
                            return (
                                f"✅ Loaded preset '{name}'",
                                preset_data.get('prompt', ''),
                                preset_data.get('neg_prompt', ''),
                                preset_data.get('steps', 28),
                                preset_data.get('cfg', 7.0),
                                preset_data.get('sampler', 'Euler')
                            )
                        
                        def delete_selected_preset(name):
                            """Delete selected preset."""
                            if not name:
                                return "❌ Please select a preset", preset_manager.list_presets()
                            
                            success = preset_manager.delete_preset(name)
                            if success:
                                return f"✅ Deleted '{name}'", preset_manager.list_presets()
                            return f"❌ Failed to delete", preset_manager.list_presets()
                        
                        preset_save_btn.click(
                            save_current_preset,
                            inputs=[
                                preset_name_input,
                                prompt_gui, neg_prompt_gui, steps_gui, cfg_gui, 
                                sampler_gui, img_width_gui, img_height_gui
                            ],
                            outputs=[preset_status_gui, preset_list_gui]
                        )
                        
                        preset_load_btn.click(
                            load_selected_preset,
                            inputs=[preset_list_gui],
                            outputs=[
                                preset_status_gui, prompt_gui, neg_prompt_gui,
                                steps_gui, cfg_gui, sampler_gui
                            ]
                        )
                        
                        preset_delete_btn.click(
                            delete_selected_preset,
                            inputs=[preset_list_gui],
                            outputs=[preset_status_gui, preset_list_gui]
                        )
                    
                    with gr.Tab("📝 Prompt Templates"):
                        gr.Markdown("""
                        ### Prompt Template System
                        Use pre-defined templates with variable substitution for faster prompting.
                        """)
                        
                        template_selector = gr.Dropdown(
                            label="Select Template",
                            choices=PromptTemplateSystem.list_templates(),
                            value=None
                        )
                        
                        template_preview = gr.Markdown(label="Template Preview")
                        
                        with gr.Accordion("Template Variables", open=False) as template_vars_accordion:
                            template_variables_gui = gr.JSON(label="Variables", value={})
                        
                        rendered_prompt_gui = gr.Textbox(
                            label="Rendered Prompt",
                            interactive=False,
                            lines=4
                        )
                        
                        use_template_btn = gr.Button("Use This Prompt", variant="primary")
                        
                        def preview_template(template_id):
                            """Preview selected template."""
                            if not template_id:
                                return "Select a template to preview", {}, ""
                            
                            template = PromptTemplateSystem.get_template(template_id)
                            if not template:
                                return "Template not found", {}, ""
                            
                            preview_md = f"""### {template['name']}
{template.get('description', '')}

**Template:**
```
{template['template']}
```

**Available Variables:**
"""
                            for var_name, desc in template.get('variables', {}).items():
                                preview_md += f"- `{var_name}`: {desc}\n"
                            
                            return preview_md, template.get('variables', {}), template['template']
                        
                        def render_template_with_vars(template_id, variables_str):
                            """Render template with user-provided variables."""
                            if not template_id:
                                return ""
                            
                            try:
                                variables = json.loads(variables_str) if variables_str else {}
                            except:
                                variables = {}
                            
                            result = PromptTemplateSystem.render_template(template_id, variables)
                            return result or "Error rendering template"
                        
                        template_selector.change(
                            preview_template,
                            inputs=[template_selector],
                            outputs=[template_preview, template_variables_gui, rendered_prompt_gui]
                        )
                        
                        use_template_btn.click(
                            lambda prompt: prompt,
                            inputs=[rendered_prompt_gui],
                            outputs=[prompt_gui]
                        )
                    
                    with gr.Tab("📊 System Monitor"):
                        gr.Markdown("""
                        ### System Information & Monitoring
                        Real-time system status and cache management.
                        """)
                        
                        refresh_system_btn = gr.Button("🔄 Refresh System Info", variant="primary")
                        system_info_display = gr.Markdown(label="System Information")
                        
                        with gr.Accordion("Cache Management", open=True):
                            cache_stats_display = gr.Markdown(label="Cache Statistics")
                            cache_cleanup_btn = gr.Button("🧹 Cleanup Cache")
                            cache_clear_old_btn = gr.Button("🗑️ Clear Old Cache (24h+)")
                        
                        gen_stats_display = gr.Markdown(label="Generation Statistics")
                        
                        def refresh_system_info():
                            """Refresh and display system information."""
                            info = get_system_info()
                            return format_system_info_md(info)
                        
                        def handle_cache_cleanup():
                            """Perform cache cleanup."""
                            removed = cache_manager.cleanup()
                            stats = cache_manager.get_cache_stats()
                            stats_md = f"""### Cache Statistics
- **Files:** {stats.get('file_count', 'N/A')}
- **Size:** {stats.get('total_size_gb', 'N/A')} GB / {stats.get('max_size_gb', 'N/A')} GB
- **Usage:** {stats.get('usage_percent', 'N/A')}%

✅ Removed {removed} old files."""
                            return stats_md
                        
                        def handle_clear_old_cache():
                            """Clear cache entries older than 24 hours."""
                            removed = cache_manager.clear_old_entries(86400)
                            return f"🗑️ Cleared {removed} old cache entries."
                        
                        refresh_system_btn.click(
                            refresh_system_info,
                            outputs=[system_info_display]
                        )
                        
                        cache_cleanup_btn.click(
                            handle_cache_cleanup,
                            outputs=[cache_stats_display]
                        )
                        
                        cache_clear_old_btn.click(
                            handle_clear_old_cache,
                            outputs=[cache_stats_display]
                        )
                        
                        # Auto-refresh on tab open
                        system_info_display.value = format_system_info_md(get_system_info())
                        cache_stats_display.value = f"""### Cache Statistics
{json.dumps(cache_manager.get_cache_stats(), indent=2)}"""

                    # End of new tabs section

                    with gr.Tab("🎨 Quick Tasks"):
                        gr.Markdown("""
                        ### Quick Task Shortcuts
                        Fast access to common Stable Diffusion tasks powered by **stablepy**.
                        
                        Select a task to auto-configure the settings:
                        """)
                        
                        with gr.Row():
                            quick_task_selector = gr.Dropdown(
                                label="Select Quick Task",
                                choices=[
                                    "📝 Text to Image (txt2img)",
                                    "🖼️ Image to Image (img2img",
                                    "🎭 Inpainting",
                                    "🌅 Outpainting (Expand Image)",
                                    "✏️ Sketch to Image (Scribble)",
                                    "📐 Pose to Image (OpenPose)",
                                    "🔲 Line Art to Image",
                                    "📏 Depth to Image",
                                    "🎨 Colorize (LineArt/Anyline)",
                                    "🔄 Style Transfer (IP-Adapter)",
                                ],
                                value="📝 Text to Image (txt2img)",
                                interactive=True
                            )
                            apply_quick_task_btn = gr.Button("✅ Apply Task Settings", variant="primary")
                        
                        quick_task_info = gr.Markdown(label="Task Info")
                        quick_task_status = gr.HTML(label="Status")
                        
                        def apply_quick_task(task_name):
                            """Apply pre-configured settings for selected task."""
                            task_configs = {
                                "📝 Text to Image (txt2img)": {
                                    "task": "txt2img",
                                    "info": "**Text to Image**: Generate images from text prompts.\n- No input image needed\n- Best for: Creating new artwork from descriptions"
                                },
                                "🖼️ Image to Image (img2img": {
                                    "task": "img2img",
                                    "info": "**Image to Image**: Transform existing images.\n- Requires input image\n- Strength controls transformation intensity"
                                },
                                "🎭 Inpainting": {
                                    "task": "inpaint",
                                    "info": "**Inpainting**: Edit or restore parts of an image.\n- Use mask to select area to modify\n- Great for: Fixing artifacts, changing elements"
                                },
                                "🌅 Outpainting (Expand Image)": {
                                    "task": "outpaint",
                                    "info": "**Outpainting**: Expand images beyond their borders.\n- Uses repaint ControlNet\n- Direction: Expand left, right, top, bottom, or all sides"
                                },
                                "✏️ Sketch to Image (Scribble)": {
                                    "task": "scribble ControlNet",
                                    "info": "**Sketch to Image**: Convert rough sketches to detailed images.\n- Upload a simple sketch/drawing\n- Works best with simple line drawings"
                                },
                                "📐 Pose to Image (OpenPose)": {
                                    "task": "openpose ControlNet",
                                    "info": "**Pose to Image**: Generate images with specific poses.\n- Uses Openpose ControlNet\n- Maintains human body structure and pose"
                                },
                                "🔲 Line Art to Image": {
                                    "task": "lineart ControlNet",
                                    "info": "**Line Art to Image**: Convert line art to full illustrations.\n- Preserves line structure\n- Adds color, shading, and details"
                                },
                                "📏 Depth to Image": {
                                    "task": "depth ControlNet",
                                    "info": "**Depth to Image**: Generate from depth maps.\n- Creates 3D-aware compositions\n- Good for: Architectural scenes, landscapes"
                                },
                                "🎨 Colorize (LineArt/Anyline)": {
                                    "task": "anyline ControlNet",
                                    "info": "**Colorize**: Add color to line drawings.\n- Uses Anyline or LineArt Standard\n- Perfect for manga/coloring book pages"
                                },
                                "🔄 Style Transfer (IP-Adapter)": {
                                    "task": "txt2img",  # IP-Adapter is additional
                                    "info": "**Style Transfer**: Apply style from reference image.\n- Enable IP-Adapter tab\n- Upload style reference image"
                                },
                            }
                            
                            config = task_configs.get(task_name, task_configs["📝 Text to Image (txt2img)"])
                            
                            return (
                                config["info"],
                                f"<b>✅ Applied:</b> {task_name}<br><b>Task set to:</b> {config['task']}",
                                gr.update(value=config["task"]),
                                gr.update(visible=config["task"] in ["inpaint", "outpaint", "img2img", 
                                                                       "scribble ControlNet", "openpose ControlNet",
                                                                       "lineart ControlNet", "depth ControlNet",
                                                                       "anyline ControlNet", "canny ControlNet",
                                                                       "mlsd ControlNet", "softedge ControlNet",
                                                                       "segmentation ControlNet", "normalbae ControlNet",
                                                                       "lineart_anime ControlNet", "tile ControlNet",
                                                                       "recolor ControlNet", "repaint ControlNet",
                                                                       "ip2p ControlNet", "pattern ControlNet",
                                                                       "shuffle ControlNet", "zoe_depth ControlNet",
                                                                       "lineart_standard ControlNet", "teed ControlNet"])
                            )
                        
                        apply_quick_task_btn.click(
                            fn=apply_quick_task,
                            inputs=[quick_task_selector],
                            outputs=[quick_task_info, quick_task_status, task_gui, menu_cn]
                        )

                    with gr.Tab("🖼️ Outpainting"):
                        gr.Markdown("""
                        ### Outpainting Tool (stablepy)
                        Expand images beyond their original borders using the **Repaint** ControlNet.
                        
                        **Features:**
                        - 🌅 Expand in any direction (left, right, top, bottom)
                        - 🎨 Seamless blending with original image
                        - ⚡ Powered by stablepy's repaint ControlNet
                        """)
                        
                        with gr.Row():
                            outpaint_input = gr.Image(label="Original Image", type="pil", sources=["upload", "clipboard"], height=300)
                            outpaint_output = gr.Image(label="Outpainted Result", height=300)
                        
                        with gr.Row():
                            outpaint_direction = gr.Radio(
                                label="Expand Direction",
                                choices=["Right", "Left", "Top", "Bottom", "All Sides"],
                                value="Right",
                                interactive=True
                            )
                            outpaint_pixels = gr.Slider(label="Pixels to Expand", minimum=64, maximum=512, step=64, value=256)
                        
                        with gr.Row():
                            outpaint_prompt = gr.Textbox(lines=3, placeholder="Describe what should appear in the expanded area...", label="Outpaint Prompt")
                            outpaint_neg_prompt = gr.Textbox(lines=2, placeholder="What to avoid in expanded area...", label="Negative Prompt")
                        
                        with gr.Accordion("Advanced Outpaint Settings", open=False):
                            with gr.Row():
                                outpaint_strength = gr.Slider(label="Denoising Strength", minimum=0.1, maximum=1.0, step=0.05, value=0.65)
                                outpaint_steps = gr.Slider(label="Steps", minimum=10, maximum=50, step=1, value=30)
                            with gr.Row():
                                outpaint_cfg = gr.Slider(label="CFG Scale", minimum=1.0, maximum=20.0, step=0.5, value=7.5)
                                outpaint_seed = gr.Number(label="Seed (-1 random)", value=-1)
                        
                        outpaint_generate_btn = gr.Button("🌅 Generate Outpaint", variant="primary")
                        outpaint_status = gr.HTML(label="Status")
                        
                        def run_outpaint(image, direction, pixels, prompt, neg_prompt, strength, steps, cfg, seed):
                            """Run outpainting using stablepy."""
                            if image is None:
                                return None, "<b>❌ Error:</b> Please upload an image first!"
                            
                            try:
                                # This would call sd_gen with outpaint/repaint task
                                # For now, show status message
                                status = f"""
                                <b>🚀 Starting Outpaint...</b><br>
                                - <b>Direction:</b> {direction}<br>
                                - <b>Pixels:</b> {pixels}<br>
                                - <b>Task:</b> repaint ControlNet<br>
                                <br>
                                <i>⚠️ Note: Ensure model is loaded before generating</i>
                                """
                                return image, status  # Placeholder - actual implementation would generate
                            except Exception as e:
                                return None, f"<b>❌ Error:</b> {str(e)}"
                        
                        outpaint_generate_btn.click(
                            fn=run_outpaint,
                            inputs=[
                                outpaint_input, outpaint_direction, outpaint_pixels,
                                outpaint_prompt, outpaint_neg_prompt,
                                outpaint_strength, outpaint_steps, outpaint_cfg, outpaint_seed
                            ],
                            outputs=[outpaint_output, outpaint_status]
                        )

                    interface_mode_gui.change(
                        fn=change_interface_mode,
                        inputs=[interface_mode_gui],
                        outputs=[
                            menu_cn, menu_inpaint, menu_lora, menu_emb, menu_hires,
                            menu_ip, menu_t2i, menu_perf, menu_ad, menu_face,
                            menu_gen, menu_model, menu_from_image, menu_negative,
                        ],
                    )

                    optimization_gui.change(
                        fn=set_optimization,
                        inputs=[optimization_gui],
                        outputs=[
                            free_u_gui, free_u_gui_perf, generator_in_cpu_gui,
                            gui_rescale_gui, pag_scale_gui,
                        ],
                    )

                    quality_selector_gui.change(
                        fn=set_quick_presets,
                        inputs=[quality_selector_gui],
                        outputs=[prompt_gui, neg_prompt_gui],
                    )

                    style_selector_gui.change(
                        fn=process_style_prompt,
                        inputs=[style_selector_gui, prompt_gui],
                        outputs=[prompt_gui],
                    )

                    sampler_selector_gui.change(
                        fn=set_sampler_settings,
                        inputs=[sampler_selector_gui],
                        outputs=[
                            steps_gui, cfg_gui, sampler_gui, schedule_type_gui,
                            schedule_prediction_type_gui,
                        ],
                    )

                    def get_t2i_model_info(model_name):
                        return get_t2i_model_info(model_name)

                    model_name_gui.change(
                        get_t2i_model_info,
                        inputs=[model_name_gui],
                        outputs=[model_info_gui],
                    )

                    def _load_model(*args):
                        yield from sd_gen.load_new_model(*args)

                    load_model_gui = gr.HTML(elem_id="load_model", elem_classes="contain")

                    load_model_gui_event = model_name_gui.change(
                        _load_model,
                        inputs=[model_name_gui, vae_model_gui, task_gui, controlnet_model_gui],
                        outputs=[load_model_gui, result_images, actual_task_info],
                    )

                    load_model_gui_event.then(
                        lambda: gr.update(visible=True),
                        None,
                        [generate_button],
                    )

                    generate_button.click(
                        fn=_load_model,
                        inputs=[model_name_gui, vae_model_gui, task_gui, controlnet_model_gui],
                        outputs=[load_model_gui, result_images, actual_task_info],
                    ).then(
                        fn=sd_gen_generate_pipeline,
                        inputs=[
                            prompt_gui, neg_prompt_gui, num_images_gui, steps_gui, cfg_gui, clip_skip_gui, seed_gui,
                            lora1_gui, lora_scale1_gui, lora2_gui, lora_scale2_gui, lora3_gui, lora_scale3_gui, lora4_gui, lora_scale4_gui, lora5_gui, lora_scale5_gui, lora6_gui, lora_scale6_gui, lora7_gui, lora_scale7_gui,
                            sampler_gui, schedule_type_gui, schedule_prediction_type_gui, img_height_gui, img_width_gui, model_name_gui, vae_model_gui, task_gui,
                            image_control_gui, preprocessor_name_gui, preprocess_resolution_gui, image_resolution_gui,
                            style_prompt_gui, style_json_gui, image_mask_gui, inpaint_strength_gui,
                            low_threshold_gui, high_threshold_gui, value_threshold_gui, distance_threshold_gui,
                            recolor_gamma_correction_gui, tile_blur_sigma_gui,
                            controlnet_output_scaling_in_unet_gui, controlnet_start_threshold_gui, controlnet_stop_threshold_gui,
                            textual_inversion_gui, syntax_weights_gui,
                            hires_upscaler_gui, hires_upscaler_increases_size_gui, hires_tile_size_gui, hires_tile_overlap_gui,
                            hires_steps_gui, hires_denoising_strength_gui, hires_sampler_gui, hires_prompt_gui, hires_neg_prompt_gui,
                            hires_before_adetailer_gui, hires_after_adetailer_gui, hires_schedule_type_gui, hires_guidance_scale_gui,
                            controlnet_model_gui,
                            loop_generation_gui, leave_progress_bar_gui, disable_progress_bar_gui,
                            image_previews_gui, display_images_gui, save_generated_images_gui,
                            filename_pattern_gui, image_storage_location_gui,
                            retain_compel_gui, retain_detailfix_gui, retain_hires_gui,
                            t2i_adapter_preprocessor_gui, t2i_adapter_conditioning_scale_gui, t2i_adapter_conditioning_factor_gui,
                            enable_live_preview_gui, free_u_gui_perf, generator_in_cpu_gui,
                            adetailer_inpaint_only_gui, adetailer_verbose_gui, adetailer_sampler_gui,
                            adetailer_active_a_gui, prompt_ad_a_gui, negative_prompt_ad_a_gui, strength_ad_a_gui,
                            face_detector_ad_a_gui, person_detector_ad_a_gui, hand_detector_ad_a_gui,
                            mask_dilation_a_gui, mask_blur_a_gui, mask_padding_a_gui,
                            adetailer_active_b_gui, prompt_ad_b_gui, negative_prompt_ad_b_gui, strength_ad_b_gui,
                            face_detector_ad_b_gui, person_detector_ad_b_gui, hand_detector_ad_b_gui,
                            mask_dilation_b_gui, mask_blur_b_gui, mask_padding_b_gui,
                            retain_task_cache_gui, guidance_rescale_gui,
                            image_ip1_gui, mask_ip1_gui, model_ip1_gui, mode_ip1_gui, scale_ip1_gui,
                            image_ip2_gui, mask_ip2_gui, model_ip2_gui, mode_ip2_gui, scale_ip2_gui,
                            pag_scale_gui, face_restoration_model_gui, face_restoration_visibility_gui, face_restoration_weight_gui,
                            load_lora_cpu_gui, verbose_info_gui, gpu_duration_gui,
                        ],
                        outputs=[load_model_gui, result_images, actual_task_info],
                    ).then(
                        fn=save_gallery_history,
                        inputs=[result_images, result_images_files],
                        outputs=[history_gallery, history_files],
                    )

                    examples = EXAMPLES_GUI if "EXAMPLES_GUI" in dir() else []
                    if examples:
                        gr.Examples(
                            examples=examples,
                            inputs=[
                                prompt_gui, neg_prompt_gui, num_images_gui, steps_gui, cfg_gui, clip_skip_gui, seed_gui,
                                lora1_gui, lora_scale1_gui, lora2_gui, lora_scale2_gui, lora3_gui, lora_scale3_gui, lora4_gui, lora_scale4_gui, lora5_gui, lora_scale5_gui, lora6_gui, lora_scale6_gui, lora7_gui, lora_scale7_gui,
                                sampler_gui, img_height_gui, img_width_gui, model_name_gui,
                            ],
                            outputs=[load_model_gui, result_images, actual_task_info],
                            cache_examples=False,
                        )

            gr.Markdown(RESOURCES)

        with gr.Tab("Inpaint mask maker", render=True):

            with gr.Row():
                with gr.Column(scale=2):
                    image_base = gr.ImageEditor(
                        sources=["upload", "clipboard"],
                        brush=gr.Brush(
                            default_size="16",
                            color_mode="fixed",
                            colors=[
                                "rgba(0, 0, 0, 1)",
                                "rgba(0, 0, 0, 0.1)",
                                "rgba(255, 255, 255, 0.1)",
                            ]
                        ),
                        eraser=gr.Eraser(default_size="16"),
                        render=True,
                        visible=False,
                        interactive=False,
                    )

                with gr.Column(scale=1):
                    mask_create_btn = gr.Button(value="Create Mask", variant="primary")
                    mask_image = gr.Image(type="numpy", label="Mask for Inpaint", interactive=False)

            mask_create_btn.click(
                fn=create_mask_now,
                inputs=image_base,
                outputs=mask_image,
            )

        with gr.Tab("Preprocessor", render=True):
            preprocessor_tab()

        with gr.Tab("Upscale", render=True):
            with gr.Row():
                with gr.Column():
                    image_upscale_input = gr.Image(label="Source Image", type="pil", sources=["upload", "clipboard"])
                    upscale_model_name = gr.Dropdown(label="Upscaler", choices=UPSCALER_KEYS, value=UPSCALER_KEYS[0])
                    upscale_factor = gr.Slider(minimum=1.0, maximum=4.0, step=0.1, value=2.0, label="Upscale factor")
                    upscale_btn = gr.Button(value="UPSCALE IMAGE", variant="primary")

                with gr.Column():
                    image_upscale_output = gr.Image(type="filepath", label="Upscaled Image", interactive=False)

            upscale_btn.click(
                fn=process_upscale,
                inputs=[image_upscale_input, upscale_model_name, upscale_factor],
                outputs=[image_upscale_output],
            )

    # ==================== PUBLIC API (UNCHANGED) ====================
    
    def generate_image(
        self,
        prompt: str,
        negative_prompt: str = "",
        num_images: int = 1,
        num_inference_steps: int = 28,
        guidance_scale: float = 7.0,
        clip_skip: int = 0,
        seed: int = -1,
        lora1: str = "", lora1_wt: float = 1.0,
        lora2: str = "", lora2_wt: float = 1.0,
        lora3: str = "", lora3_wt: float = 1.0,
        lora4: str = "", lora4_wt: float = 1.0,
        lora5: str = "", lora5_wt: float = 1.0,
        lora6: str = "", lora6_wt: float = 1.0,
        lora7: str = "", lora7_wt: float = 1.0,
        sampler: str = "Euler",
        schedule_type: str = "Automatic",
        schedule_prediction_type: str = "Automatic",
        height: int = 1024,
        width: int = 1024,
        model_name: str = "votepurchase/animagine-xl-3.1",
        vae_model: str = "None",
        task: str = "txt2img",
        image_control_dict: Optional[dict] = None,
        preprocessor_name: str = "Canny",
        preprocess_resolution: int = 512,
        image_resolution: int = 1024,
        style_prompt: Optional[List[str]] = None,
        style_json: Optional[dict] = None,
        image_mask: Optional[Any] = None,
        strength: float = 0.55,
        low_threshold: int = 100,
        high_threshold: int = 200,
        value_threshold: float = 0.1,
        distance_threshold: float = 0.1,
        recolor_gamma_correction: float = 1.0,
        tile_blur_sigma: int = 9,
        control_net_output_scaling: float = 1.0,
        control_net_start_threshold: float = 0.0,
        control_net_stop_threshold: float = 1.0,
        textual_inversion: bool = False,
        prompt_syntax: str = "Classic",
        upscaler_model_path: Optional[str] = None,
        upscaler_increases_size: float = 1.2,
        upscaler_tile_size: int = 0,
        upscaler_tile_overlap: int = 8,
        hires_steps: int = 30,
        hires_denoising_strength: float = 0.55,
        hires_sampler: str = "Use same sampler",
        hires_prompt: str = "",
        hires_negative_prompt: str = "",
        adetailer_inpaint_only: bool = True,
        adetailer_verbose: bool = False,
        hires_schedule_type: str = "Use same schedule type",
        hires_guidance_scale: float = -1.0,
        controlnet_model: str = "Automatic",
        loop_generation: bool = False,
        leave_progress_bar: bool = False,
        disable_progress_bar: bool = False,
        image_previews: bool = True,
        display_images: bool = True,
        save_generated_images: bool = True,
        filename_pattern: str = "model,seed",
        image_storage_location: str = "./images/",
        retain_compel_previous_load: bool = True,
        retain_detailfix_model_previous_load: bool = True,
        retain_hires_model_previous_load: bool = True,
        t2i_adapter_preprocessor: Optional[str] = None,
        t2i_adapter_conditioning_scale: float = 0.55,
        t2i_adapter_conditioning_factor: float = 1.0,
        xformers_memory_efficient_attention: bool = True,
        free_u: bool = False,
        generator_in_cpu: bool = False,
        adetailer_sampler: str = "Use same sampler",
        adetailer_active_a: bool = False,
        prompt_ad_a: str = "",
        negative_prompt_ad_a: str = "",
        strength_ad_a: float = 0.35,
        face_detector_ad_a: bool = False,
        person_detector_ad_a: bool = True,
        hand_detector_ad_a: bool = False,
        mask_dilation_a: int = 4,
        mask_blur_a: int = 4,
        mask_padding_a: int = 32,
        adetailer_active_b: bool = False,
        prompt_ad_b: str = "",
        negative_prompt_ad_b: str = "",
        strength_ad_b: float = 0.35,
        face_detector_ad_b: bool = False,
        person_detector_ad_b: bool = True,
        hand_detector_ad_b: bool = False,
        mask_dilation_b: int = 4,
        mask_blur_b: int = 4,
        mask_padding_b: int = 32,
        cache_compel_texts: bool = True,
        guidance_rescale: float = 0.0,
        image_ip1_dict: Optional[dict] = None, mask_ip1: Optional[Any] = None,
        model_ip1: str = "plus_face", mode_ip1: str = "original", scale_ip1: float = 0.7,
        image_ip2_dict: Optional[dict] = None, mask_ip2: Optional[Any] = None,
        model_ip2: str = "base", mode_ip2: str = "style", scale_ip2: float = 0.7,
        pag_scale: float = 0.0,
        face_restoration_model: Optional[str] = None,
        face_restoration_visibility: float = 1.0,
        face_restoration_weight: float = 0.5,
        load_lora_cpu: bool = False,
        verbose_info_gui: int = 0,
        gpu_duration: int = 20,
    ) -> Tuple[str, Optional[List[str]], Optional[str]]:
        """Generate image with explicit arguments (non-streaming API)."""
        # Ensure the correct model is loaded before generation.
        _load_model(model_name, vae_model, task, controlnet_model)

        # Build argv in the exact order expected by sd_gen_generate_pipeline(*argv).
        argv: List[Any] = [
            prompt, negative_prompt, num_images, num_inference_steps, guidance_scale, clip_skip, seed,
            lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt,
            lora6, lora6_wt, lora7, lora7_wt,
            sampler, schedule_type, schedule_prediction_type,
            height, width, model_name, vae_model, task,
            image_control_dict, preprocessor_name, preprocess_resolution, image_resolution,
            style_prompt, style_json, image_mask,
            strength, low_threshold, high_threshold, value_threshold, distance_threshold,
            recolor_gamma_correction, tile_blur_sigma,
            control_net_output_scaling, control_net_start_threshold, control_net_stop_threshold,
            textual_inversion, prompt_syntax,
            upscaler_model_path, upscaler_increases_size, upscaler_tile_size, upscaler_tile_overlap,
            hires_steps, hires_denoising_strength, hires_sampler, hires_prompt, hires_negative_prompt,
            adetailer_inpaint_only, adetailer_verbose, hires_schedule_type, hires_guidance_scale,
            controlnet_model,
            loop_generation, leave_progress_bar, disable_progress_bar,
            image_previews, display_images, save_generated_images,
            filename_pattern, image_storage_location,
            retain_compel_previous_load, retain_detailfix_model_previous_load, retain_hires_model_previous_load,
            t2i_adapter_preprocessor, t2i_adapter_conditioning_scale, t2i_adapter_conditioning_factor,
            xformers_memory_efficient_attention, free_u, generator_in_cpu,
            adetailer_sampler,
            adetailer_active_a, prompt_ad_a, negative_prompt_ad_a, strength_ad_a,
            face_detector_ad_a, person_detector_ad_a, hand_detector_ad_a,
            mask_dilation_a, mask_blur_a, mask_padding_a,
            adetailer_active_b, prompt_ad_b, negative_prompt_ad_b, strength_ad_b,
            face_detector_ad_b, person_detector_ad_b, hand_detector_ad_b,
            mask_dilation_b, mask_blur_b, mask_padding_b,
            cache_compel_texts, guidance_rescale,
            image_ip1_dict, mask_ip1, model_ip1, mode_ip1, scale_ip1,
            image_ip2_dict, mask_ip2, model_ip2, mode_ip2, scale_ip2,
            pag_scale, face_restoration_model, face_restoration_visibility, face_restoration_weight,
            load_lora_cpu, verbose_info_gui, gpu_duration,
        ]

        last: Tuple[str, Optional[List[str]], Optional[str]] = ("COMPLETE", None, None)
        for last in _generate_image(argv):
            pass
        return last

    # Streaming API
    def generate_image_stream(
        self,
        # Same signature as generate_image; kept duplicated for clarity and API docs.
        prompt: str,
        negative_prompt: str = "",
        num_images: int = 1,
        num_inference_steps: int = 28,
        guidance_scale: float = 7.0,
        clip_skip: int = 0,
        seed: int = -1,
        lora1: str = "", lora1_wt: float = 1.0,
        lora2: str = "", lora2_wt: float = 1.0,
        lora3: str = "", lora3_wt: float = 1.0,
        lora4: str = "", lora4_wt: float = 1.0,
        lora5: str = "", lora5_wt: float = 1.0,
        lora6: str = "", lora6_wt: float = 1.0,
        lora7: str = "", lora7_wt: float = 1.0,
        sampler: str = "Euler",
        schedule_type: str = "Automatic",
        schedule_prediction_type: str = "Automatic",
        height: int = 1024,
        width: int = 1024,
        model_name: str = "votepurchase/animagine-xl-3.1",
        vae_model: str = "None",
        task: str = "txt2img",
        image_control_dict: Optional[dict] = None,
        preprocessor_name: str = "Canny",
        preprocess_resolution: int = 512,
        image_resolution: int = 1024,
        style_prompt: Optional[List[str]] = None,
        style_json: Optional[dict] = None,
        image_mask: Optional[Any] = None,
        strength: float = 0.55,
        low_threshold: int = 100,
        high_threshold: int = 200,
        value_threshold: float = 0.1,
        distance_threshold: float = 0.1,
        recolor_gamma_correction: float = 1.0,
        tile_blur_sigma: int = 9,
        control_net_output_scaling: float = 1.0,
        control_net_start_threshold: float = 0.0,
        control_net_stop_threshold: float = 1.0,
        textual_inversion: bool = False,
        prompt_syntax: str = "Classic",
        upscaler_model_path: Optional[str] = None,
        upscaler_increases_size: float = 1.2,
        upscaler_tile_size: int = 0,
        upscaler_tile_overlap: int = 8,
        hires_steps: int = 30,
        hires_denoising_strength: float = 0.55,
        hires_sampler: str = "Use same sampler",
        hires_prompt: str = "",
        hires_negative_prompt: str = "",
        adetailer_inpaint_only: bool = True,
        adetailer_verbose: bool = False,
        hires_schedule_type: str = "Use same schedule type",
        hires_guidance_scale: float = -1.0,
        controlnet_model: str = "Automatic",
        loop_generation: bool = False,
        leave_progress_bar: bool = False,
        disable_progress_bar: bool = False,
        image_previews: bool = True,
        display_images: bool = True,
        save_generated_images: bool = True,
        filename_pattern: str = "model,seed",
        image_storage_location: str = "./images/",
        retain_compel_previous_load: bool = True,
        retain_detailfix_model_previous_load: bool = True,
        retain_hires_model_previous_load: bool = True,
        t2i_adapter_preprocessor: Optional[str] = None,
        t2i_adapter_conditioning_scale: float = 0.55,
        t2i_adapter_conditioning_factor: float = 1.0,
        xformers_memory_efficient_attention: bool = True,
        free_u: bool = False,
        generator_in_cpu: bool = False,
        adetailer_sampler: str = "Use same sampler",
        adetailer_active_a: bool = False,
        prompt_ad_a: str = "",
        negative_prompt_ad_a: str = "",
        strength_ad_a: float = 0.35,
        face_detector_ad_a: bool = False,
        person_detector_ad_a: bool = True,
        hand_detector_ad_a: bool = False,
        mask_dilation_a: int = 4,
        mask_blur_a: int = 4,
        mask_padding_a: int = 32,
        adetailer_active_b: bool = False,
        prompt_ad_b: str = "",
        negative_prompt_ad_b: str = "",
        strength_ad_b: float = 0.35,
        face_detector_ad_b: bool = False,
        person_detector_ad_b: bool = True,
        hand_detector_ad_b: bool = False,
        mask_dilation_b: int = 4,
        mask_blur_b: int = 4,
        mask_padding_b: int = 32,
        cache_compel_texts: bool = True,
        guidance_rescale: float = 0.0,
        image_ip1_dict: Optional[dict] = None, mask_ip1: Optional[Any] = None,
        model_ip1: str = "plus_face", mode_ip1: str = "original", scale_ip1: float = 0.7,
        image_ip2_dict: Optional[dict] = None, mask_ip2: Optional[Any] = None,
        model_ip2: str = "base", mode_ip2: str = "style", scale_ip2: float = 0.7,
        pag_scale: float = 0.0,
        face_restoration_model: Optional[str] = None,
        face_restoration_visibility: float = 1.0,
        face_restoration_weight: float = 0.5,
        load_lora_cpu: bool = False,
        verbose_info_gui: int = 0,
        gpu_duration: int = 20,
    ) -> Generator[Tuple[str, Optional[List[str]], Optional[str]], None, None]:
        """Generate image with streaming updates."""
        _load_model(model_name, vae_model, task, controlnet_model)

        argv: List[Any] = [
            prompt, negative_prompt, num_images, num_inference_steps, guidance_scale, clip_skip, seed,
            lora1, lora1_wt, lora2, lora2_wt, lora3, lora3_wt, lora4, lora4_wt, lora5, lora5_wt,
            lora6, lora6_wt, lora7, lora7_wt,
            sampler, schedule_type, schedule_prediction_type,
            height, width, model_name, vae_model, task,
            image_control_dict, preprocessor_name, preprocess_resolution, image_resolution,
            style_prompt, style_json, image_mask,
            strength, low_threshold, high_threshold, value_threshold, distance_threshold,
            recolor_gamma_correction, tile_blur_sigma,
            control_net_output_scaling, control_net_start_threshold, control_net_stop_threshold,
            textual_inversion, prompt_syntax,
            upscaler_model_path, upscaler_increases_size, upscaler_tile_size, upscaler_tile_overlap,
            hires_steps, hires_denoising_strength, hires_sampler, hires_prompt, hires_negative_prompt,
            adetailer_inpaint_only, adetailer_verbose, hires_schedule_type, hires_guidance_scale,
            controlnet_model,
            loop_generation, leave_progress_bar, disable_progress_bar,
            image_previews, display_images, save_generated_images,
            filename_pattern, image_storage_location,
            retain_compel_previous_load, retain_detailfix_model_previous_load, retain_hires_model_previous_load,
            t2i_adapter_preprocessor, t2i_adapter_conditioning_scale, t2i_adapter_conditioning_factor,
            xformers_memory_efficient_attention, free_u, generator_in_cpu,
            adetailer_sampler,
            adetailer_active_a, prompt_ad_a, negative_prompt_ad_a, strength_ad_a,
            face_detector_ad_a, person_detector_ad_a, hand_detector_ad_a,
            mask_dilation_a, mask_blur_a, mask_padding_a,
            adetailer_active_b, prompt_ad_b, negative_prompt_ad_b, strength_ad_b,
            face_detector_ad_b, person_detector_ad_b, hand_detector_ad_b,
            mask_dilation_b, mask_blur_b, mask_padding_b,
            cache_compel_texts, guidance_rescale,
            image_ip1_dict, mask_ip1, model_ip1, mode_ip1, scale_ip1,
            image_ip2_dict, mask_ip2, model_ip2, mode_ip2, scale_ip2,
            pag_scale, face_restoration_model, face_restoration_visibility, face_restoration_weight,
            load_lora_cpu, verbose_info_gui, gpu_duration,
        ]

        yield from _generate_image(argv)

    # Register APIs (wrapped in try-except for Gradio version compatibility)
    try:
        gr.api(generate_image, api_name="generate_image", api_visibility="public", queue=True, concurrency_id="gpu")
        gr.api(generate_image_stream, api_name="generate_image_stream", api_visibility="public", queue=True, concurrency_id="gpu")
    except (ValueError, TypeError) as e:
        print(f"[Warning] API endpoint registration skipped: {e}")

    gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires GPU Space)")


if __name__ == "__main__":
    app.queue()
    app.launch(
        show_error=True,
        share=args.share_enabled,
        debug=True,
        ssr_mode=args.ssr,
        mcp_server=False,
        allowed_paths=[allowed_path],
        theme=args.theme,
        css=CSS,
    )