app.py

#!/usr/bin/env python3
"""
Pro Realism Edit Studio - Enhanced Edition
=========================================

Advanced image editing and enhancement studio powered by:
- Qwen-Image-Edit-2511 with Phr00t's Rapid-AIO v23 accelerated transformer
- Nomos-family tiled upscaling
- Qwen-aware masked cleanup and final photographic detail

Author: Enhanced with Hugging Face CLI and image generation expertise
Version: 1.0.4
"""

import gradio as gr
import numpy as np
import random
import torch
import spaces
import os
import time
import tempfile
from pathlib import Path

def _configure_model_cache_environment():
    bucket_root = os.environ.get("MODEL_BUCKET_DIR", "/models").strip()
    fallback_root = os.environ.get("MODEL_CACHE_FALLBACK_DIR", "/data/models").strip()
    candidates = [bucket_root, fallback_root, "/tmp/pro-realism-models"]
    cache_root = None

    for candidate in candidates:
        if not candidate:
            continue
        path = Path(candidate)
        if path.exists() or candidate.startswith(("/data", "/tmp")):
            try:
                path.mkdir(parents=True, exist_ok=True)
                probe = path / ".write-test"
                probe.write_text("ok", encoding="utf-8")
                probe.unlink(missing_ok=True)
                cache_root = path
                break
            except Exception:
                continue

    if cache_root is None:
        return None

    hf_home = cache_root / "huggingface"
    hf_hub_cache = hf_home / "hub"
    torch_home = cache_root / "torch"
    ultralytics_cache = cache_root / "ultralytics"

    for path in (hf_home, hf_hub_cache, torch_home, ultralytics_cache):
        path.mkdir(parents=True, exist_ok=True)

    os.environ.setdefault("HF_HOME", str(hf_home))
    os.environ.setdefault("HF_HUB_CACHE", str(hf_hub_cache))
    os.environ.setdefault("TRANSFORMERS_CACHE", str(hf_hub_cache))
    os.environ.setdefault("DIFFUSERS_CACHE", str(hf_hub_cache))
    os.environ.setdefault("TORCH_HOME", str(torch_home))
    os.environ.setdefault("ULTRALYTICS_CACHE_DIR", str(ultralytics_cache))
    os.environ.setdefault("YOLO_CONFIG_DIR", str(ultralytics_cache))
    print(f"Model cache root: {cache_root}")
    return cache_root

MODEL_CACHE_ROOT = _configure_model_cache_environment()

# Advanced imports
from accelerate import init_empty_weights
from collections import OrderedDict
from PIL import Image, ImageChops, ImageEnhance, ImageFilter, ImageOps
from diffusers.models import QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel
from diffusers.models.model_loading_utils import load_model_dict_into_meta
from huggingface_hub import hf_hub_download, HfApi, login, whoami
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
from safetensors import safe_open

from gradio_client import Client, handle_file

# ============================================================================
# CONFIGURATION - Model IDs and Parameters
# ============================================================================

# Base model configuration
BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
APP_VERSION = "1.0.4"
PHR00T_REPO_ID = os.environ.get("PHR00T_REPO_ID", "Phr00t/Qwen-Image-Edit-Rapid-AIO").strip()
RAPID_TRANSFORMER_FILENAME = os.environ.get(
    "RAPID_TRANSFORMER_FILENAME",
    "v23/Qwen-Rapid-AIO-NSFW-v23.safetensors",
).strip()
PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."
VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()

# Qwen-optimized enhancement configuration
UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "Phips/4xNomos8k_atd_jpg").strip()
UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "4xNomos8k_atd_jpg.safetensors").strip()
UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "64"))
ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "2048"))
ENHANCE_QWEN_DOWNSCALE_TRIGGER_EDGE = int(os.environ.get("ENHANCE_QWEN_DOWNSCALE_TRIGGER_EDGE", "1536"))
ENHANCE_QWEN_DOWNSCALE_FACTOR = float(os.environ.get("ENHANCE_QWEN_DOWNSCALE_FACTOR", "0.75"))
ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.010"))
ENHANCE_DETAILER_ENABLED = os.environ.get("ENHANCE_DETAILER_ENABLED", "true").lower() == "true"
ENHANCE_DETAILER_MAX_REGIONS = int(os.environ.get("ENHANCE_DETAILER_MAX_REGIONS", "8"))
ENHANCE_DETAILER_MIN_REGION_AREA = float(os.environ.get("ENHANCE_DETAILER_MIN_REGION_AREA", "0.003"))
ENHANCE_YOLO_SAM_ENABLED = os.environ.get("ENHANCE_YOLO_SAM_ENABLED", "true").lower() == "true"
ENHANCE_YOLO_MODEL = os.environ.get("ENHANCE_YOLO_MODEL", "yolo26n-seg.pt").strip()
ENHANCE_SAM_MODEL = os.environ.get("ENHANCE_SAM_MODEL", "sam_l.pt").strip()
ENHANCE_YOLO_CONF = float(os.environ.get("ENHANCE_YOLO_CONF", "0.18"))
ENHANCE_DA3_PRESERVE_ENABLED = os.environ.get("ENHANCE_DA3_PRESERVE_ENABLED", "false").lower() == "true"
ENHANCE_DA3_MODEL_ID = os.environ.get("ENHANCE_DA3_MODEL_ID", "depth-anything/DA3-BASE").strip()

# Advanced Detail Enhancement Configuration
DETAIL_ENHANCEMENT_ENABLED = os.environ.get("DETAIL_ENHANCEMENT_ENABLED", "true").lower() == "true"
SMART_SHARPENING_STRENGTH = float(os.environ.get("SMART_SHARPENING_STRENGTH", "1.15"))

# ============================================================================
# ENHANCEMENT MODES
# ============================================================================

ENHANCE_MODE_OFF = "Off"
ENHANCE_MODE_UPSCALE = "Upscale"
ENHANCE_MODE_CLEAN = "Clean"
ENHANCE_MODE_MAX_DETAIL = "Max Detail"
ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]

# ============================================================================
# GLOBAL MODEL CACHE
# ============================================================================

_upscaler_model = None
_detail_enhancement_model = None
_yolo_model = None
_sam_model = None
_da3_model = None

# ============================================================================
# HUGGING FACE CLI EXPERT FUNCTIONS
# ============================================================================

def download_model_with_retry(repo_id, filename, max_retries=3):
    """Download model with retry logic and error handling"""
    for attempt in range(max_retries):
        try:
            return hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                cache_dir=os.environ.get("HF_HUB_CACHE"),
            )
        except Exception as e:
            if attempt == max_retries - 1:
                raise RuntimeError(f"Failed to download {filename} from {repo_id} after {max_retries} attempts: {e}")
            time.sleep(2 ** attempt)
    return None

# ============================================================================
# VIDEO GENERATION (Preserved from original)
# ============================================================================

def turn_into_video(input_image, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
    """Convert image edit into video transition"""
    if not VIDEO_SPACE_ID:
        raise gr.Error("Video generation is not configured for this Space.")
    if not input_image or not output_images:
        raise gr.Error("Please generate an output image first.")

    progress(0.02, desc="Preparing images...")

    def extract_pil(img_entry):
        if isinstance(img_entry, tuple) and isinstance(img_entry[0], Image.Image):
            return img_entry[0]
        elif isinstance(img_entry, Image.Image):
            return img_entry
        elif isinstance(img_entry, str):
            return Image.open(img_entry)
        else:
            raise gr.Error(f"Unsupported image format: {type(img_entry)}")

    start_img = extract_pil(input_image)
    end_img = extract_pil(output_images[0])

    progress(0.10, desc="Saving temp files...")

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
        start_img.save(tmp_start.name)
        end_img.save(tmp_end.name)

    progress(0.20, desc="Connecting to video Space...")

    client = Client(VIDEO_SPACE_ID)

    progress(0.35, desc="Generating video...")

    video_path, seed = client.predict(
        start_image_pil=handle_file(tmp_start.name),
        end_image_pil=handle_file(tmp_end.name),
        prompt=prompt or "smooth cinematic transition",
        api_name="/generate_video"
    )

    progress(0.95, desc="Finalizing...")
    return video_path['video']


# ============================================================================
# HISTORY MANAGEMENT (Enhanced)
# ============================================================================

def update_history(new_images, history):
    """Updates the history gallery with the new images."""
    time.sleep(0.3)
    if history is None:
        history = []
    if new_images is not None and len(new_images) > 0:
        if not isinstance(history, list):
            history = list(history) if history else []
        for img in new_images:
            history.insert(0, img)
    history = history[:50]
    return history

def use_history_as_input(evt: gr.SelectData):
    """Sets the selected history image into the Image 1 slot."""
    if evt.value is not None:
        return gr.update(value=evt.value)
    return gr.update()

# ============================================================================
# MODEL LOADING (Enhanced with better error handling)
# ============================================================================

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_phr00t_rapid_transformer(torch_dtype):
    """Load Phr00t's Rapid-AIO v23 transformer with enhanced error handling"""
    checkpoint_path = download_model_with_retry(PHR00T_REPO_ID, RAPID_TRANSFORMER_FILENAME)
    
    try:
        config = DiffusersQwenImageTransformer2DModel.load_config(
            BASE_MODEL_ID,
            subfolder="transformer",
            cache_dir=os.environ.get("HF_HUB_CACHE"),
        )
    except Exception as e:
        raise RuntimeError(f"Failed to load config for {BASE_MODEL_ID}: {e}")
    
    with init_empty_weights():
        transformer = DiffusersQwenImageTransformer2DModel.from_config(config)

    expected_keys = set(transformer.state_dict().keys())
    state_dict = OrderedDict()
    
    try:
        with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint:
            for key in checkpoint.keys():
                if not key.startswith(PHR00T_TRANSFORMER_PREFIX):
                    continue
                mapped_key = key.removeprefix(PHR00T_TRANSFORMER_PREFIX)
                if mapped_key in expected_keys:
                    state_dict[mapped_key] = checkpoint.get_tensor(key)
    except Exception as e:
        raise RuntimeError(f"Failed to load checkpoint from {checkpoint_path}: {e}")

    missing_keys = sorted(expected_keys.difference(state_dict.keys()))
    if missing_keys:
        sample = ", ".join(missing_keys[:20])
        raise RuntimeError(
            f"Phr00t Rapid-AIO transformer checkpoint is missing {len(missing_keys)} "
            f"required diffusers keys after prefix conversion. First missing keys: {sample}"
        )

    try:
        load_model_dict_into_meta(transformer, state_dict, dtype=torch_dtype)
    except Exception as e:
        raise RuntimeError(f"Failed to load state dict into meta: {e}")
    
    meta_parameters = [name for name, parameter in transformer.named_parameters() if parameter.is_meta]
    if meta_parameters:
        sample = ", ".join(meta_parameters[:20])
        raise RuntimeError(
            f"Phr00t Rapid-AIO transformer still has {len(meta_parameters)} meta parameters "
            f"after loading. First meta parameters: {sample}"
        )

    transformer.eval()
    return transformer

# Load main pipeline
try:
    pipe = QwenImageEditPlusPipeline.from_pretrained(
        BASE_MODEL_ID,
        transformer=load_phr00t_rapid_transformer(dtype),
        torch_dtype=dtype,
        cache_dir=os.environ.get("HF_HUB_CACHE"),
    ).to(device)
    print("Successfully loaded Qwen-Image-Edit-2511 with Rapid-AIO v23 transformer")
except Exception as e:
    print(f"Failed to load main pipeline: {e}")
    raise

# Apply optimizations
pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
print("Applied FA3 attention processor optimization")

# ============================================================================
# QWEN-OPTIMIZED UPSCALER
# ============================================================================

def load_upscaler_model():
    """Load a Spandrel upscaler, defaulting to the Nomos model used by this Space."""
    global _upscaler_model
    if _upscaler_model is not None:
        return _upscaler_model

    try:
        import spandrel
        import spandrel_extra_arches
        spandrel_extra_arches.install()
    except ImportError as exc:
        raise gr.Error("Enhance mode requires spandrel and spandrel_extra_arches to be installed. "
                       "Install with: pip install spandrel spandrel_extra_arches") from exc

    candidates = [
        (UPSCALER_MODEL_ID, UPSCALER_MODEL_FILENAME),
        ("Phips/4xNomos8k_atd_jpg", "4xNomos8k_atd_jpg.safetensors"),
    ]
    seen = set()
    errors = []
    for repo_id, filename in candidates:
        key = (repo_id, filename)
        if key in seen:
            continue
        seen.add(key)
        try:
            model_path = download_model_with_retry(repo_id, filename)
            model = spandrel.ModelLoader().load_from_file(model_path)
            model.eval().to(device)
            _upscaler_model = model
            print(f"Successfully loaded upscaler: {repo_id}/{filename}")
            return _upscaler_model
        except Exception as exc:
            errors.append(f"{repo_id}/{filename}: {exc}")
            print(f"Failed to load upscaler {repo_id}/{filename}: {exc}")

    raise gr.Error(f"Failed to load all upscaler models: {' | '.join(errors)}")

def image_to_tensor(image):
    """Convert PIL Image to tensor"""
    array = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
    return tensor.to(device)

def tensor_to_image(tensor):
    """Convert tensor to PIL Image"""
    array = tensor.squeeze(0).detach().float().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
    return Image.fromarray((array * 255.0).round().astype(np.uint8), mode="RGB")

def validate_enhance_input_size(image):
    """Validate image size for enhancement"""
    max_edge = max(image.size)
    if max_edge > ENHANCE_MAX_INPUT_EDGE:
        raise gr.Error(
            f"Enhance mode accepts images up to {ENHANCE_MAX_INPUT_EDGE}px on the longest edge. "
            f"Current image is {image.width}x{image.height}. "
            f"Consider resizing your image first."
        )

def load_yolo_model():
    global _yolo_model
    if _yolo_model is not None:
        return _yolo_model
    try:
        from ultralytics import YOLO
        _yolo_model = YOLO(ENHANCE_YOLO_MODEL)
        return _yolo_model
    except Exception as exc:
        print(f"YOLO detailer unavailable: {exc}")
        return None

def load_sam_model():
    global _sam_model
    if _sam_model is not None:
        return _sam_model
    try:
        from ultralytics import SAM
        _sam_model = SAM(ENHANCE_SAM_MODEL)
        return _sam_model
    except Exception as exc:
        print(f"SAM detailer unavailable: {exc}")
        return None

def load_da3_model():
    global _da3_model
    if _da3_model is not None:
        return _da3_model
    if not ENHANCE_DA3_PRESERVE_ENABLED:
        return None
    try:
        from depth_anything_3.api import DepthAnything3
        _da3_model = DepthAnything3.from_pretrained(ENHANCE_DA3_MODEL_ID).to(device)
        return _da3_model
    except Exception as exc:
        print(f"DA-3 preservation unavailable: {exc}")
        return None

def _tile_weight(height, width, overlap_y, overlap_x, touches_top, touches_bottom, touches_left, touches_right, dtype, device):
    weight = torch.ones((1, 1, height, width), dtype=dtype, device=device)
    if overlap_y > 1 and not touches_top:
        ramp = torch.linspace(0.0, 1.0, overlap_y, dtype=dtype, device=device).view(1, 1, overlap_y, 1)
        weight[:, :, :overlap_y, :] *= ramp
    if overlap_y > 1 and not touches_bottom:
        ramp = torch.linspace(1.0, 0.0, overlap_y, dtype=dtype, device=device).view(1, 1, overlap_y, 1)
        weight[:, :, -overlap_y:, :] *= ramp
    if overlap_x > 1 and not touches_left:
        ramp = torch.linspace(0.0, 1.0, overlap_x, dtype=dtype, device=device).view(1, 1, 1, overlap_x)
        weight[:, :, :, :overlap_x] *= ramp
    if overlap_x > 1 and not touches_right:
        ramp = torch.linspace(1.0, 0.0, overlap_x, dtype=dtype, device=device).view(1, 1, 1, overlap_x)
        weight[:, :, :, -overlap_x:] *= ramp
    return weight

def tile_upscale(image, scale=4):
    """Tiled Spandrel upscaling with feathered overlaps to avoid visible seams."""
    validate_enhance_input_size(image)
    model = load_upscaler_model()
    tensor = image_to_tensor(image)
    _, _, height, width = tensor.shape
    
    base_tile_size = UPSCALER_TILE_SIZE
    optimal_tile_size = min(base_tile_size, max(height, width) // 2)
    tile_size = max(64, optimal_tile_size)
    overlap = max(0, min(UPSCALER_TILE_OVERLAP, tile_size // 2))
    step = max(1, tile_size - overlap)
    
    y_positions = list(range(0, height, step))
    if y_positions[-1] + tile_size < height:
        y_positions.append(max(0, height - tile_size))
    
    x_positions = list(range(0, width, step))
    if x_positions[-1] + tile_size < width:
        x_positions.append(max(0, width - tile_size))
    
    y_positions = sorted(set(y_positions))
    x_positions = sorted(set(x_positions))
    
    output = None
    weights = None

    with torch.inference_mode():
        for y in y_positions:
            for x in x_positions:
                y1 = min(y + tile_size, height)
                x1 = min(x + tile_size, width)
                tile = tensor[:, :, y:y1, x:x1]
                
                upscaled_tile = model(tile)
                if isinstance(upscaled_tile, (tuple, list)):
                    upscaled_tile = upscaled_tile[0]
                upscaled_tile = upscaled_tile.clamp(0, 1)
                
                scale_y = upscaled_tile.shape[-2] // tile.shape[-2]
                scale_x = upscaled_tile.shape[-1] // tile.shape[-1]
                
                if output is None:
                    output = torch.zeros(
                        (1, 3, height * scale_y, width * scale_x),
                        dtype=upscaled_tile.dtype,
                        device=upscaled_tile.device,
                    )
                    weights = torch.zeros_like(output)
                
                oy0, oy1 = y * scale_y, y1 * scale_y
                ox0, ox1 = x * scale_x, x1 * scale_x
                blend = _tile_weight(
                    upscaled_tile.shape[-2],
                    upscaled_tile.shape[-1],
                    min(overlap * scale_y, max(1, upscaled_tile.shape[-2] // 2)),
                    min(overlap * scale_x, max(1, upscaled_tile.shape[-1] // 2)),
                    y == 0,
                    y1 == height,
                    x == 0,
                    x1 == width,
                    upscaled_tile.dtype,
                    upscaled_tile.device,
                )
                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile * blend
                weights[:, :, oy0:oy1, ox0:ox1] += blend

    output = output / weights.clamp_min(1)
    return tensor_to_image(output)

def advanced_tile_upscale(image, scale=4):
    return tile_upscale(image, scale=scale)

# ============================================================================
# ENHANCED DETAILER (Multi-stage processing)
# ============================================================================

def smart_sharpen(image, strength=1.15):
    """
    Smart sharpening with edge detection to avoid oversharpening smooth areas
    """
    if strength <= 0:
        return image
    
    img_array = np.array(image.convert("RGB"))
    
    if strength > 1.0:
        enhanced = ImageEnhance.Sharpness(image).enhance(strength)
        
        gray = image.convert("L")
        edges = gray.filter(ImageFilter.FIND_EDGES)
        edge_mask = edges.filter(ImageFilter.GaussianBlur(radius=1))
        edge_mask = edge_mask.point(lambda x: min(x * 0.3, 255))
        
        sharpened_array = np.array(enhanced)
        original_array = img_array
        edge_array = np.array(edge_mask).astype(float) / 255.0
        
        for c in range(3):
            sharpened_array[:, :, c] = (
                edge_array * sharpened_array[:, :, c] + 
                (1 - edge_array) * original_array[:, :, c]
            )
        
        image = Image.fromarray(np.clip(sharpened_array, 0, 255).astype(np.uint8))
    
    return image

def add_ultra_detail(image, strength=0.8):
    """
    Add ultra-fine details using high-frequency enhancement
    """
    if strength <= 0:
        return image
    
    original = image.convert("RGB")
    blurred = original.filter(ImageFilter.GaussianBlur(radius=2))
    
    high_freq = ImageChops.subtract(original, blurred)
    high_freq_enhanced = ImageEnhance.Contrast(high_freq).enhance(1.0 + strength)
    result = ImageChops.add(original, high_freq_enhanced)
    
    return result

def apply_high_frequency_details(image, amount=0.6):
    """
    Apply high-frequency detail enhancement for crisp textures
    """
    if amount <= 0:
        return image
    
    scales = [1, 2, 4]
    result = image.convert("RGB")
    
    for scale in scales:
        blurred = result.filter(ImageFilter.GaussianBlur(radius=scale))
        high_freq = ImageChops.subtract(result, blurred)
        enhanced_hf = ImageEnhance.Contrast(high_freq).enhance(1.0 + amount * 0.3)
        result = ImageChops.add(result, enhanced_hf)
    
    return result

# ============================================================================
# LEGACY CLEANUP HELPERS
# ============================================================================

def remove_artifacts(image):
    """Remove compression artifacts and noise"""
    denoised = image.filter(ImageFilter.MedianFilter(size=3))
    smoothed = denoised.filter(ImageFilter.GaussianBlur(radius=0.5))
    result = Image.blend(image, smoothed, alpha=0.3)
    return result

def enhanced_skin_repair(image):
    """Enhanced skin repair with better color detection and blending"""
    base = image.convert("RGB")
    
    ycbcr = np.asarray(base.convert("YCbCr"))
    y, cb, cr = ycbcr[:, :, 0], ycbcr[:, :, 1], ycbcr[:, :, 2]
    
    skin_mask = (
        (cr > 130) & (cr < 170) &
        (cb > 70) & (cb < 140) &
        (y > 80)
    ).astype(np.uint8) * 255
    
    try:
        import cv2
        kernel = np.ones((5, 5), np.uint8)
        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
        skin_mask = cv2.GaussianBlur(skin_mask, (7, 7), 0)
    except ImportError:
        from scipy import ndimage
        skin_mask = ndimage.binary_opening(skin_mask > 128, structure=np.ones((3, 3))).astype(np.uint8) * 255
        skin_mask = ndimage.gaussian_filter(skin_mask, sigma=3)
    
    mask_image = Image.fromarray(skin_mask, mode="L")
    repaired = base.filter(ImageFilter.MedianFilter(size=3))
    repaired = repaired.filter(ImageFilter.GaussianBlur(radius=0.4))
    non_skin = ImageOps.invert(mask_image)
    sharpened = ImageEnhance.Sharpness(base).enhance(1.15)
    blended = Image.composite(repaired, sharpened, mask_image)
    result = ImageEnhance.Sharpness(blended).enhance(1.05)
    return result

# Original skin repair functions (preserved for compatibility)
def skin_repair_mask(image):
    ycbcr = np.asarray(image.convert("YCbCr"))
    cb = ycbcr[:, :, 1]
    cr = ycbcr[:, :, 2]
    mask = (
        (cr >= 135)
        & (cr <= 180)
        & (cb >= 75)
        & (cb <= 135)
    ).astype(np.uint8) * 255
    mask_image = Image.fromarray(mask, mode="L")
    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.2))

def repair_skin_texture(image):
    base = image.convert("RGB")
    mask = skin_repair_mask(base)
    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.35))
    blended = Image.composite(repaired, base, mask)
    return ImageEnhance.Sharpness(blended).enhance(1.08)

def add_film_grain(image, seed):
    base = image.convert("RGB")
    array = np.asarray(base).astype(np.float32)
    rng = np.random.default_rng(seed)
    grain = rng.normal(0.0, 255.0 * ENHANCE_GRAIN_STRENGTH, size=(array.shape[0], array.shape[1], 1))
    array = np.clip(array + grain, 0, 255)
    return Image.fromarray(array.astype(np.uint8), mode="RGB")

def _resize_to_even_dimensions(width, height):
    return max(2, width - (width % 2)), max(2, height - (height % 2))

def qwen_artifact_precondition(image):
    """
    Qwen edit outputs can show halftone/plastic texture at larger dimensions.
    A small Lanczos downsample before the SR pass suppresses that pattern while
    preserving prompt structure for the upscaler.
    """
    base = image.convert("RGB")
    long_edge = max(base.size)
    if long_edge <= ENHANCE_QWEN_DOWNSCALE_TRIGGER_EDGE:
        return base

    factor = min(0.95, max(0.50, ENHANCE_QWEN_DOWNSCALE_FACTOR))
    new_width, new_height = _resize_to_even_dimensions(
        int(round(base.width * factor)),
        int(round(base.height * factor)),
    )
    if new_width >= base.width or new_height >= base.height:
        return base
    return base.resize((new_width, new_height), Image.Resampling.LANCZOS)

def qwen_skin_mask(image):
    base = image.convert("RGB")
    ycbcr = np.asarray(base.convert("YCbCr"))
    y = ycbcr[:, :, 0]
    cb = ycbcr[:, :, 1]
    cr = ycbcr[:, :, 2]
    mask = (
        (y > 45)
        & (cb >= 72)
        & (cb <= 145)
        & (cr >= 128)
        & (cr <= 182)
    ).astype(np.uint8) * 255

    try:
        import cv2
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.GaussianBlur(mask, (0, 0), 1.6)
    except ImportError:
        mask_image = Image.fromarray(mask, mode="L")
        mask_image = mask_image.filter(ImageFilter.MinFilter(size=3))
        mask_image = mask_image.filter(ImageFilter.MaxFilter(size=5))
        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=1.4))
        mask = np.asarray(mask_image)

    return Image.fromarray(mask.astype(np.uint8), mode="L")

def qwen_defect_mask(image):
    base = image.convert("RGB")
    skin = np.asarray(qwen_skin_mask(base)).astype(np.float32) / 255.0
    gray = np.asarray(base.convert("L")).astype(np.float32)
    local = np.asarray(base.convert("L").filter(ImageFilter.MedianFilter(size=5))).astype(np.float32)
    pits = np.maximum(local - gray, 0.0)
    smears = np.abs(gray - local)
    mask = ((pits > 14.0) | (smears > 22.0)) & (skin > 0.12)
    mask = (mask.astype(np.uint8) * 255)
    mask_image = Image.fromarray(mask, mode="L")
    mask_image = mask_image.filter(ImageFilter.MaxFilter(size=3))
    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.1))

def qwen_hair_mask(image):
    base = image.convert("RGB")
    rgb = np.asarray(base).astype(np.float32)
    gray_image = base.convert("L")
    gray = np.asarray(gray_image).astype(np.float32)
    skin = np.asarray(qwen_skin_mask(base)).astype(np.float32) / 255.0
    edges = np.asarray(gray_image.filter(ImageFilter.FIND_EDGES).filter(ImageFilter.GaussianBlur(radius=0.7))).astype(np.float32)
    chroma = rgb.max(axis=2) - rgb.min(axis=2)

    dark_strands = (gray < 122) & (edges > 8) & (skin < 0.45)
    light_strands = (gray < 235) & (edges > 18) & (chroma > 8) & (skin < 0.28)
    mask = (dark_strands | light_strands).astype(np.uint8) * 255

    try:
        import cv2
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        mask = cv2.GaussianBlur(mask, (0, 0), 1.0)
    except ImportError:
        mask_image = Image.fromarray(mask, mode="L")
        mask_image = mask_image.filter(ImageFilter.MaxFilter(size=3))
        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=1.0))
        mask = np.asarray(mask_image)

    return Image.fromarray(mask.astype(np.uint8), mode="L")

def qwen_face_feature_mask(image):
    base = image.convert("RGB")
    skin = np.asarray(qwen_skin_mask(base)).astype(np.float32) / 255.0
    gray_image = base.convert("L")
    edges = np.asarray(gray_image.filter(ImageFilter.FIND_EDGES)).astype(np.float32)
    features = ((skin > 0.08) & (edges > 10)).astype(np.uint8) * 255
    defects = np.asarray(qwen_defect_mask(base)).astype(np.uint8)
    mask = np.maximum(features, defects)
    mask_image = Image.fromarray(mask, mode="L")
    mask_image = mask_image.filter(ImageFilter.MaxFilter(size=5))
    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.6))

def _results_to_person_boxes(results, image_size):
    boxes = []
    for result in results or []:
        if getattr(result, "boxes", None) is None:
            continue
        xyxy = result.boxes.xyxy.detach().cpu().numpy() if result.boxes.xyxy is not None else []
        cls = result.boxes.cls.detach().cpu().numpy() if result.boxes.cls is not None else []
        conf = result.boxes.conf.detach().cpu().numpy() if result.boxes.conf is not None else []
        for idx, box in enumerate(xyxy):
            class_id = int(cls[idx]) if idx < len(cls) else -1
            score = float(conf[idx]) if idx < len(conf) else 1.0
            if class_id == 0 and score >= ENHANCE_YOLO_CONF:
                x0, y0, x1, y1 = [int(round(value)) for value in box]
                x0 = max(0, min(image_size[0] - 1, x0))
                x1 = max(1, min(image_size[0], x1))
                y0 = max(0, min(image_size[1] - 1, y0))
                y1 = max(1, min(image_size[1], y1))
                if x1 > x0 and y1 > y0:
                    boxes.append((x0, y0, x1, y1))
    boxes.sort(key=lambda item: (item[2] - item[0]) * (item[3] - item[1]), reverse=True)
    return boxes[:ENHANCE_DETAILER_MAX_REGIONS]

def yolo_person_boxes(image):
    if not ENHANCE_YOLO_SAM_ENABLED:
        return []
    model = load_yolo_model()
    if model is None:
        return []
    try:
        results = model.predict(
            source=np.asarray(image.convert("RGB")),
            classes=[0],
            conf=ENHANCE_YOLO_CONF,
            verbose=False,
        )
        return _results_to_person_boxes(results, image.size)
    except Exception as exc:
        print(f"YOLO person detection failed: {exc}")
        return []

def sam_mask_from_boxes(image, boxes):
    if not boxes or not ENHANCE_YOLO_SAM_ENABLED:
        return Image.new("L", image.size, 0)
    model = load_sam_model()
    if model is None:
        return Image.new("L", image.size, 0)
    try:
        union = np.zeros((image.height, image.width), dtype=np.uint8)
        for box in boxes:
            results = model(np.asarray(image.convert("RGB")), bboxes=list(box), verbose=False)
            for result in results or []:
                masks = getattr(result, "masks", None)
                if masks is None or masks.data is None:
                    continue
                data = masks.data.detach().float().cpu().numpy()
                for mask in data:
                    mask_image = Image.fromarray((mask > 0.5).astype(np.uint8) * 255, mode="L")
                    mask_image = mask_image.resize(image.size, Image.Resampling.BILINEAR)
                    union = np.maximum(union, np.asarray(mask_image, dtype=np.uint8))
        return Image.fromarray(union, mode="L").filter(ImageFilter.GaussianBlur(radius=1.2))
    except Exception as exc:
        print(f"SAM box segmentation failed: {exc}")
        return Image.new("L", image.size, 0)

def da3_preservation_mask(image, person_mask):
    if not ENHANCE_DA3_PRESERVE_ENABLED:
        return person_mask
    model = load_da3_model()
    if model is None:
        return person_mask
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            prediction = model.inference(
                [np.asarray(image.convert("RGB"))],
                export_dir=tmpdir,
                export_format="npz",
            )
        if isinstance(prediction, dict):
            depth = prediction.get("depth")
            if depth is None:
                depth = prediction.get("depths")
        else:
            depth = getattr(prediction, "depth", None)
            if depth is None:
                depth = getattr(prediction, "depths", None)
        if depth is None:
            return person_mask
        depth_array = np.asarray(depth[0] if isinstance(depth, (list, tuple)) else depth).astype(np.float32)
        if depth_array.ndim > 2:
            depth_array = depth_array.squeeze()
        depth_array -= depth_array.min()
        depth_array /= max(float(depth_array.max()), 1e-6)
        depth_image = Image.fromarray((depth_array * 255).astype(np.uint8), mode="L").resize(image.size, Image.Resampling.BILINEAR)
        foreground = depth_image.point(lambda value: 255 if value >= 32 else 0).filter(ImageFilter.GaussianBlur(radius=1.4))
        return ImageChops.multiply(person_mask.convert("L"), foreground)
    except Exception as exc:
        print(f"DA-3 preservation mask failed: {exc}")
        return person_mask

def yolo_sam_person_mask_and_boxes(image):
    boxes = yolo_person_boxes(image)
    person_mask = sam_mask_from_boxes(image, boxes)
    if person_mask.getbbox() is None:
        return person_mask, boxes
    return da3_preservation_mask(image, person_mask), boxes

def _mask_to_boxes(mask_image, max_regions=ENHANCE_DETAILER_MAX_REGIONS):
    mask = np.asarray(mask_image.convert("L"))
    binary = (mask > 24).astype(np.uint8)
    min_area = max(32, int(binary.shape[0] * binary.shape[1] * ENHANCE_DETAILER_MIN_REGION_AREA))
    boxes = []

    try:
        import cv2
        count, labels, stats, _ = cv2.connectedComponentsWithStats(binary, connectivity=8)
        for label in range(1, count):
            x, y, width, height, area = stats[label]
            if area >= min_area:
                boxes.append((int(x), int(y), int(x + width), int(y + height), int(area)))
    except ImportError:
        bbox = mask_image.point(lambda value: 255 if value > 24 else 0).getbbox()
        if bbox:
            x0, y0, x1, y1 = bbox
            boxes.append((x0, y0, x1, y1, (x1 - x0) * (y1 - y0)))

    boxes.sort(key=lambda item: item[4], reverse=True)
    return [box[:4] for box in boxes[:max_regions]]

def _expand_box(box, image_size, pad_ratio=0.18, min_size=192):
    x0, y0, x1, y1 = box
    width = x1 - x0
    height = y1 - y0
    pad = int(max(width, height) * pad_ratio)
    if width < min_size:
        extra = (min_size - width) // 2
        x0 -= extra
        x1 += extra
    if height < min_size:
        extra = (min_size - height) // 2
        y0 -= extra
        y1 += extra
    return (
        max(0, x0 - pad),
        max(0, y0 - pad),
        min(image_size[0], x1 + pad),
        min(image_size[1], y1 + pad),
    )

def _local_detail_crop(crop, mask_crop, strength=0.35, hair=False):
    base = crop.convert("RGB")
    mask = mask_crop.convert("L").filter(ImageFilter.GaussianBlur(radius=1.8))
    strength = max(0.0, min(1.0, strength))

    if hair:
        detailed = base.filter(ImageFilter.UnsharpMask(radius=0.55, percent=int(130 + 120 * strength), threshold=2))
        detailed = ImageEnhance.Contrast(detailed).enhance(1.0 + 0.08 * strength)
        edge_mask = base.convert("L").filter(ImageFilter.FIND_EDGES).filter(ImageFilter.GaussianBlur(radius=0.7))
        mask = ImageChops.multiply(mask, edge_mask.point(lambda value: min(255, int(value * 2.2))))
    else:
        repaired = masked_texture_repair(base, strength=0.35 + 0.30 * strength)
        detailed = repaired.filter(ImageFilter.UnsharpMask(radius=0.75, percent=int(80 + 90 * strength), threshold=3))
        detailed = ImageEnhance.Contrast(detailed).enhance(1.0 + 0.045 * strength)

    mask = mask.point(lambda value: int(value * strength))
    return Image.composite(detailed, base, mask)

def qwen_tiled_detailer_pass(image, strength=0.35, include_hair=True):
    if not ENHANCE_DETAILER_ENABLED or strength <= 0:
        return image.convert("RGB")

    result = image.convert("RGB")
    person_mask, person_boxes = yolo_sam_person_mask_and_boxes(result)
    person_mask = person_mask.convert("L")

    face_mask = qwen_face_feature_mask(result)
    if person_mask.getbbox() is not None:
        face_mask = ImageChops.multiply(face_mask, person_mask)

    region_specs = [(face_mask, strength, False, 0.22, person_boxes)]
    if include_hair:
        hair_mask = qwen_hair_mask(result)
        if person_mask.getbbox() is not None:
            hair_mask = ImageChops.multiply(hair_mask, person_mask)
        region_specs.append((hair_mask, strength * 0.85, True, 0.12, person_boxes))

    for mask, region_strength, hair, pad_ratio, boxes in region_specs:
        region_boxes = boxes or _mask_to_boxes(mask)
        for box in region_boxes:
            expanded = _expand_box(box, result.size, pad_ratio=pad_ratio, min_size=192)
            crop = result.crop(expanded)
            mask_crop = mask.crop(expanded)
            if mask_crop.getbbox() is None:
                mask_crop = Image.new("L", crop.size, 255)
            detailed = _local_detail_crop(crop, mask_crop, strength=region_strength, hair=hair)
            result.paste(detailed, expanded, mask_crop.filter(ImageFilter.GaussianBlur(radius=2.5)))

    return result

def masked_texture_repair(image, strength=0.55):
    base = image.convert("RGB")
    mask = qwen_defect_mask(base)
    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.22))
    repaired = ImageEnhance.Sharpness(repaired).enhance(1.06)
    mask = mask.point(lambda value: int(value * max(0.0, min(1.0, strength))))
    return Image.composite(repaired, base, mask)

def qwen_micro_detail(image, amount=0.35):
    if amount <= 0:
        return image.convert("RGB")
    base = image.convert("RGB")
    sharpened = base.filter(ImageFilter.UnsharpMask(radius=0.9, percent=int(95 * amount), threshold=3))
    contrast = ImageEnhance.Contrast(sharpened).enhance(1.0 + amount * 0.05)
    return Image.blend(base, contrast, alpha=min(0.65, amount))

def add_photographic_grain(image, seed):
    return add_film_grain(image, seed)

# ============================================================================
# ENHANCED APPLY ENHANCEMENT (Main enhancement pipeline)
# ============================================================================

def apply_enhancement(image, enhance_mode, seed=0, progress=None):
    """Apply the Qwen 2511 / Rapid-AIO v23 post-process pipeline."""
    mode = enhance_mode or ENHANCE_MODE_OFF
    if mode not in ENHANCE_MODE_CHOICES:
        raise gr.Error(f"Unknown enhance mode: {mode}")
    if mode == ENHANCE_MODE_OFF:
        return image

    enhanced = image.convert("RGB")
    total_steps = {
        ENHANCE_MODE_UPSCALE: 2,
        ENHANCE_MODE_CLEAN: 4,
        ENHANCE_MODE_MAX_DETAIL: 7,
    }[mode]
    step = 0

    if progress:
        step += 1
        progress(0.85 * step / total_steps, desc="Preparing Qwen output...")
    enhanced = qwen_artifact_precondition(enhanced)

    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL):
        if progress:
            step += 1
            progress(0.85 * step / total_steps, desc="Repairing Qwen skin artifacts...")
        enhanced = masked_texture_repair(enhanced, strength=0.55 if mode == ENHANCE_MODE_CLEAN else 0.68)
        enhanced = qwen_micro_detail(enhanced, amount=0.18)

        if progress:
            step += 1
            progress(0.85 * step / total_steps, desc="Detailing face regions...")
        enhanced = qwen_tiled_detailer_pass(
            enhanced,
            strength=0.32 if mode == ENHANCE_MODE_CLEAN else 0.48,
            include_hair=(mode == ENHANCE_MODE_MAX_DETAIL),
        )

    if mode == ENHANCE_MODE_MAX_DETAIL:
        if progress:
            step += 1
            progress(0.85 * step / total_steps, desc="Building micro detail...")
        enhanced = qwen_micro_detail(enhanced, amount=0.36)

    if progress:
        step += 1
        progress(0.85 * step / total_steps, desc="Upscaling with Nomos tiles...")
    enhanced = tile_upscale(enhanced)

    if mode == ENHANCE_MODE_MAX_DETAIL:
        if progress:
            step += 1
            progress(0.94, desc="Detailing final face and hair tiles...")
        enhanced = qwen_tiled_detailer_pass(enhanced, strength=0.30, include_hair=True)
        if progress:
            step += 1
            progress(0.98, desc="Adding final photographic grain...")
        enhanced = qwen_micro_detail(enhanced, amount=0.22)
        enhanced = add_photographic_grain(enhanced, seed)

    return enhanced

# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

def use_output_as_input(output_images):
    """Move the first output image into the Image 1 slot."""
    if not output_images:
        return gr.update()
    first = output_images[0]
    path = first[0] if isinstance(first, (list, tuple)) else first
    return gr.update(value=path)

def check_gpu_memory():
    """Check available GPU memory"""
    if device == "cuda":
        try:
            total = torch.cuda.get_device_properties(0).total_memory
            reserved = torch.cuda.memory_reserved(0)
            allocated = torch.cuda.memory_allocated(0)
            free = total - reserved
            print(f"GPU Memory: Total={total/1024**3:.2f}GB, Reserved={reserved/1024**3:.2f}GB, Allocated={allocated/1024**3:.2f}GB, Free={free/1024**3:.2f}GB")
            return free > 1024**3
        except Exception as e:
            print(f"Failed to check GPU memory: {e}")
            return True
    return True

def clear_gpu_cache():
    """Clear GPU cache to free up memory"""
    if device == "cuda":
        try:
            torch.cuda.empty_cache()
            import gc
            gc.collect()
            print("GPU cache cleared")
        except Exception as e:
            print(f"Failed to clear GPU cache: {e}")

# ============================================================================
# MAIN INFERENCE FUNCTION (Enhanced)
# ============================================================================

MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU(duration=60)
def infer(
    image_1,
    image_2,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    enhance_mode=ENHANCE_MODE_OFF,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Enhanced image generation with advanced editing and enhancement options
    """
    negative_prompt = " "

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    pil_images = []
    for img in (image_1, image_2):
        if img is None:
            continue
        try:
            if isinstance(img, str):
                pil_images.append(Image.open(img).convert("RGB"))
            elif isinstance(img, Image.Image):
                pil_images.append(img.convert("RGB"))
            elif hasattr(img, "name"):
                pil_images.append(Image.open(img.name).convert("RGB"))
        except Exception:
            continue

    if height==256 and width==256:
        height, width = None, None
    
    print(f"Generation Parameters:")
    print(f"   Prompt: '{prompt}'")
    print(f"   Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
    print(f"   Size: {width}x{height}, Enhance Mode: {enhance_mode}")

    if not check_gpu_memory():
        clear_gpu_cache()
        if not check_gpu_memory():
            raise gr.Error("Insufficient GPU memory. Please reduce image size or close other applications.")

    try:
        images_pil = pipe(
            image=pil_images if len(pil_images) > 0 else None,
            prompt=prompt,
            height=height,
            width=width,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=num_images_per_prompt,
        ).images
    except Exception as e:
        clear_gpu_cache()
        raise gr.Error(f"Image generation failed: {e}")

    if enhance_mode != ENHANCE_MODE_OFF:
        images_pil = [
            apply_enhancement(img, enhance_mode, seed=seed + idx, progress=progress)
            for idx, img in enumerate(images_pil)
        ]

    output_paths = []
    os.makedirs("outputs", exist_ok=True)
    for idx, img in enumerate(images_pil):
        output_path = f"outputs/output_{seed}_{idx}_{int(time.time()*1000)}.png"
        img.save(output_path)
        output_paths.append(output_path)

    clear_gpu_cache()

    return output_paths, seed, gr.update(visible=True), gr.update(visible=bool(VIDEO_SPACE_ID))


# ============================================================================
# UI LAYOUT (Enhanced)
# ============================================================================

css = """
#col-container {
    margin: 0 auto;
    max-width: 1024px;
}
#logo-title {
    text-align: center;
}
#logo-title h1 {
    margin-bottom: 0;
}
#logo-title h2 {
    color: #5b47d1;
    font-style: italic;
    margin-top: 0;
}
#edit_text{margin-top: -62px !important}
.enhance-info {
    font-size: 0.9em;
    color: #666;
    margin-top: 5px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(f"""
        <!-- v{APP_VERSION} -->
        <div id="logo-title">
            <h1>Pro Realism Edit Studio - Enhanced</h1>
            <h2>Rapid Edit with Qwen-aware Nomos enhancement</h2>
        </div>
        """)
        
        gr.Markdown("""
        Powered by:
        - [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) 
        - [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer
        - YOLO-boxed SAM person masks with Nomos-family 4x tiled upscaling
        
        Upload an image and enter your prompt to edit it.
        
        Pro Tips:
        - Use Clean for portraits with smudged or pitted skin
        - Use Max Detail for hair, eyes, lashes, texture, grain, and sharper final output
        """)
        
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
                    image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)

                prompt = gr.Text(
                    label="Prompt",
                    show_label=True,
                    placeholder="Enter your prompt here...",
                )
                
                enhance_mode = gr.Radio(
                    label="Enhance Mode",
                    choices=ENHANCE_MODE_CHOICES,
                    value=ENHANCE_MODE_OFF,
                    interactive=True,
                    info="Choose enhancement level for your output"
                )
                
                enhance_info = gr.Markdown("""
                Enhancement Options:
                - Off: No post-processing
                - Upscale: Qwen precondition + 4x Nomos tiled upscale
                - Clean: YOLO/SAM person mask + tiled face detail + upscaling
                - Max Detail: YOLO/SAM person mask + tiled face/hair detail + micro detail + upscaling + final grain
                """, visible=False)
                
                run_button = gr.Button("Generate!", variant="primary")
                
                with gr.Accordion("Advanced Settings", open=False):
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )

                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                    with gr.Row():
                        true_guidance_scale = gr.Slider(
                            label="True guidance scale",
                            minimum=1.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.0
                        )

                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=40,
                            step=1,
                            value=4,
                        )
                    
                    with gr.Row():
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=2048,
                            step=8,
                            value=None,
                        )
                        
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=2048,
                            step=8,
                            value=None,
                        )
                    
                    gr.Markdown("""
                    Performance Tips:
                    - Use 4 steps for fastest results
                    - Increase steps (8-20) for better quality
                    - Lower guidance scale for more creative freedom
                    """)

            with gr.Column():
                result = gr.Gallery(label="Result", show_label=False, type="filepath")
                with gr.Row():
                    use_output_btn = gr.Button("Use as input", variant="secondary", size="sm", visible=False)
                    turn_video_btn = gr.Button("Turn into Video", variant="secondary", size="sm", visible=False)
                output_video = gr.Video(label="Generated Video", autoplay=True, visible=False)

                with gr.Row():
                    gr.Markdown("### History")
                    clear_history_button = gr.Button("Clear History", size="sm", variant="stop")
                
                history_gallery = gr.Gallery(
                    label="Click any image to use as input", 
                    interactive=False,
                    show_label=True,
                    visible=True
                )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            image_1,
            image_2,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            height,
            width,
            enhance_mode,
        ],
        outputs=[result, seed, use_output_btn, turn_video_btn],

    ).then(
        fn=update_history,
        inputs=[result, history_gallery],
        outputs=history_gallery,
    )

    enhance_mode.change(
        fn=lambda mode: gr.update(visible=mode != ENHANCE_MODE_OFF),
        inputs=[enhance_mode],
        outputs=[enhance_info]
    )

    use_output_btn.click(
        fn=use_output_as_input,
        inputs=[result],
        outputs=[image_1]
    )

    history_gallery.select(
        fn=use_history_as_input,
        inputs=None,
        outputs=[image_1],
    )
    
    clear_history_button.click(
        fn=lambda: [],
        inputs=None,
        outputs=history_gallery,
    )

    turn_video_btn.click(
        fn=lambda: gr.update(visible=True),   
        inputs=None,
        outputs=[output_video],
    ).then(
        fn=turn_into_video,
        inputs=[image_1, result, prompt],
        outputs=[output_video],
    )


if __name__ == "__main__":
    print(f"Device: {device}")
    if device == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    check_gpu_memory()
    print(f"Starting Pro Realism Edit Studio v{APP_VERSION}")
    demo.launch()