#!/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
- Real-ESRGAN for high-quality upscaling
- GFPGAN/CodeFormer for face restoration
- Multi-stage detail enhancement pipeline

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

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

# Advanced imports
from accelerate import init_empty_weights
from collections import OrderedDict
from PIL import Image, 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.0"
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()

# Enhanced Upscaler Configuration
UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "ai-forever/Real-ESRGAN").strip()
UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "RealESRGAN_x4plus.pth").strip()
UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "64"))  # Increased overlap for better blending
ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "2048"))  # Increased from 1280
ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.015"))  # Reduced from 0.018

# Face Restoration Configuration
FACE_RESTORATION_MODEL = os.environ.get("FACE_RESTORATION_MODEL", "Xintao/GFPGAN").strip()
FACE_RESTORATION_WEIGHTS = os.environ.get("FACE_RESTORATION_WEIGHTS", "GFPGANv1.3.pth").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 Only"
ENHANCE_MODE_CLEAN = "Clean & Restore"
ENHANCE_MODE_MAX_DETAIL = "Max Detail"
ENHANCE_MODE_FACE_ENHANCE = "Face Enhance"
ENHANCE_MODE_FULL_ENHANCE = "Full Enhance"
ENHANCE_MODE_CHOICES = [
    ENHANCE_MODE_OFF, 
    ENHANCE_MODE_UPSCALE, 
    ENHANCE_MODE_CLEAN, 
    ENHANCE_MODE_MAX_DETAIL,
    ENHANCE_MODE_FACE_ENHANCE,
    ENHANCE_MODE_FULL_ENHANCE
]

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

_upscaler_model = None
_face_restoration_model = None
_detail_enhancement_model = None

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

def check_hf_login():
    """Check if user is logged in to Hugging Face Hub"""
    try:
        return whoami() is not None
    except Exception:
        return False

def ensure_hf_login():
    """Ensure user is logged in, prompt if not"""
    if not check_hf_login():
        try:
            login()
            return True
        except Exception as e:
            print(f"Hugging Face login failed: {e}")
            return False
    return True

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)
        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)  # Exponential backoff
    return None

def get_model_info(repo_id):
    """Get model information from Hugging Face Hub"""
    try:
        api = HfApi()
        model_info = api.model_info(repo_id)
        return model_info
    except Exception as e:
        print(f"Failed to get model info for {repo_id}: {e}")
        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)  # Reduced delay for better responsiveness
    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]  # Increased from 20 to 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",
        )
    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
    ).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")

# ============================================================================
# ENHANCED UPSCALER (Real-ESRGAN based)
# ============================================================================

def load_upscaler_model():
    """Load Real-ESRGAN model for high-quality upscaling"""
    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

    try:
        model_path = download_model_with_retry(UPSCALER_MODEL_ID, UPSCALER_MODEL_FILENAME)
        model = spandrel.ModelLoader().load_from_file(model_path)
        model.eval().to(device)
        _upscaler_model = model
        print(f"✅ Successfully loaded upscaler: {UPSCALER_MODEL_ID}/{UPSCALER_MODEL_FILENAME}")
        return _upscaler_model
    except Exception as e:
        print(f"❌ Failed to load upscaler model: {e}")
        # Fallback to original Nomos model
        print("🔄 Falling back to Nomos upscaler...")
        try:
            model_path = download_model_with_retry("Phips/4xNomos8k_atd_jpg", "4xNomos8k_atd_jpg.safetensors")
            model = spandrel.ModelLoader().load_from_file(model_path)
            model.eval().to(device)
            _upscaler_model = model
            return _upscaler_model
        except Exception as fallback_error:
            raise gr.Error(f"Failed to load all upscaler models: {e} | {fallback_error}")

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 advanced_tile_upscale(image, scale=4):
    """
    Advanced tiling upscaler with improved blending and edge handling
    Uses Real-ESRGAN for superior quality compared to Nomos
    """
    validate_enhance_input_size(image)
    model = load_upscaler_model()
    tensor = image_to_tensor(image)
    _, _, height, width = tensor.shape
    
    # Adaptive tile size based on image dimensions
    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)
    
    # Ensure full coverage with edge tiles
    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]
                
                # Process tile through upscaler
                upscaled_tile = model(tile).clamp(0, 1)
                
                # Calculate scale factors
                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
                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile
                weights[:, :, oy0:oy1, ox0:ox1] += 1

    # Normalize overlapping regions
    output = output / weights.clamp_min(1)
    return tensor_to_image(output)

# ============================================================================
# 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
    
    # Convert to array for processing
    img_array = np.array(image.convert("RGB"))
    
    # Apply adaptive sharpening
    if strength > 1.0:
        # Use ImageEnhance for basic sharpening
        enhanced = ImageEnhance.Sharpness(image).enhance(strength)
        
        # Additional edge-aware sharpening
        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))  # Normalize edge strength
        
        # Blend sharpened version with original based on edge strength
        sharpened_array = np.array(enhanced)
        original_array = img_array
        edge_array = np.array(edge_mask).astype(float) / 255.0
        
        # Create edge-aware blend
        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
    
    # Apply high-pass filtering for detail extraction
    original = image.convert("RGB")
    blurred = original.filter(ImageFilter.GaussianBlur(radius=2))
    
    # Extract high-frequency details
    high_freq = ImageChops.subtract(original, blurred)
    
    # Enhance the high-frequency component
    high_freq_enhanced = ImageEnhance.Contrast(high_freq).enhance(1.0 + strength)
    
    # Add enhanced details back to original
    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
    
    # Multiple scales of detail enhancement
    scales = [1, 2, 4]  # Different blur radii for multi-scale details
    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

# ============================================================================
# ENHANCED CLEANER (Face Restoration + Artifact Removal)
# ============================================================================

def load_face_restoration_model():
    """Load GFPGAN model for face restoration"""
    global _face_restoration_model
    if _face_restoration_model is not None:
        return _face_restoration_model
    
    try:
        # Try to import face restoration libraries
        import gfpgan
        from gfpgan import GFPGANer
        
        # Download and load model
        model_path = download_model_with_retry(FACE_RESTORATION_MODEL, FACE_RESTORATION_WEIGHTS)
        
        # Initialize GFPGANer
        restorer = GFPGANer(
            model_path=model_path,
            upscale=1,  # We handle upscaling separately
            arch='clean',
            channel_multiplier=2,
            bg_upsampler=None
        )
        
        _face_restoration_model = restorer
        print("✅ Successfully loaded GFPGAN face restoration model")
        return _face_restoration_model
        
    except ImportError:
        print("⚠️  GFPGAN not available, face restoration will use fallback methods")
        return None
    except Exception as e:
        print(f"❌ Failed to load face restoration model: {e}")
        return None

def detect_faces(image):
    """Detect faces in an image and return bounding boxes"""
    try:
        import cv2
        import numpy as np
        
        # Convert PIL to numpy array
        img_array = np.array(image.convert("RGB"))
        gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
        
        # Load face cascade
        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
        
        return faces
    except ImportError:
        print("⚠️  OpenCV not available, using simple face detection fallback")
        # Simple fallback: assume center of image for portrait
        width, height = image.size
        if width > height:  # Landscape
            return []
        else:  # Portrait
            face_size = min(width, height) // 2
            x = (width - face_size) // 2
            y = (height - face_size) // 2
            return [[x, y, face_size, face_size]]
    except Exception as e:
        print(f"⚠️  Face detection failed: {e}")
        return []

def restore_faces(image):
    """Restore faces in an image using GFPGAN"""
    restorer = load_face_restoration_model()
    if restorer is None:
        print("⚠️  Face restoration model not available, using skin repair fallback")
        return repair_skin_texture(image)
    
    try:
        # Convert to numpy array
        img_array = np.array(image.convert("RGB"))
        
        # Restore faces
        restored_array, _ = restorer.enhance(img_array, has_aligned=False, only_center_face=False, paste_back=True)
        
        # Convert back to PIL
        restored_image = Image.fromarray(restored_array.astype(np.uint8))
        
        return restored_image
    except Exception as e:
        print(f"⚠️  Face restoration failed: {e}, using skin repair fallback")
        return repair_skin_texture(image)

def remove_artifacts(image):
    """Remove compression artifacts and noise"""
    # Apply mild median filtering for noise reduction
    denoised = image.filter(ImageFilter.MedianFilter(size=3))
    
    # Apply slight Gaussian blur to smooth artifacts
    smoothed = denoised.filter(ImageFilter.GaussianBlur(radius=0.5))
    
    # Blend with original to preserve details
    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")
    
    # Improved skin detection using YCbCr with better thresholds
    ycbcr = np.asarray(base.convert("YCbCr"))
    y, cb, cr = ycbcr[:, :, 0], ycbcr[:, :, 1], ycbcr[:, :, 2]
    
    # More sophisticated skin detection
    skin_mask = (
        (cr > 130) & (cr < 170) &
        (cb > 70) & (cb < 140) &
        (y > 80)  # Exclude dark areas
    ).astype(np.uint8) * 255
    
    # Apply morphological operations to clean up mask
    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:
        # Fallback without OpenCV
        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")
    
    # Apply more sophisticated skin repair
    repaired = base.filter(ImageFilter.MedianFilter(size=3))
    repaired = repaired.filter(ImageFilter.GaussianBlur(radius=0.4))
    
    # Apply selective sharpening to non-skin areas
    non_skin = ImageOps.invert(mask_image)
    sharpened = ImageEnhance.Sharpness(base).enhance(1.15)
    
    # Blend repaired skin with sharpened non-skin areas
    blended = Image.composite(repaired, sharpened, mask_image)
    
    # Final enhancement
    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")

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

def apply_enhancement(image, enhance_mode, seed=0, progress=None):
    """
    Apply various enhancement modes to the image
    
    Modes:
    - Off: No enhancement
    - Upscale Only: Just upscale the image
    - Clean & Restore: Remove artifacts, repair skin, restore faces
    - Max Detail: Full enhancement with detail boost
    - Face Enhance: Focus on face restoration
    - Full Enhance: Complete enhancement 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")
    
    # Progress tracking
    total_steps = 0
    if mode == ENHANCE_MODE_UPSCALE:
        total_steps = 1
    elif mode == ENHANCE_MODE_CLEAN:
        total_steps = 3
    elif mode == ENHANCE_MODE_MAX_DETAIL:
        total_steps = 4
    elif mode == ENHANCE_MODE_FACE_ENHANCE:
        total_steps = 2
    elif mode == ENHANCE_MODE_FULL_ENHANCE:
        total_steps = 5
    
    step = 0
    
    # Face Enhance Mode
    if mode == ENHANCE_MODE_FACE_ENHANCE:
        if progress:
            step += 1
            progress(0.5 * step / total_steps, desc="Restoring faces...")
        enhanced = restore_faces(enhanced)
        
        if progress:
            step += 1
            progress(0.5 * step / total_steps, desc="Upscaling...")
        enhanced = advanced_tile_upscale(enhanced)
        
        return enhanced
    
    # Clean & Restore Mode
    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_FULL_ENHANCE):
        if progress:
            step += 1
            progress(0.7 * step / total_steps, desc="Removing artifacts...")
        enhanced = remove_artifacts(enhanced)
        
        if progress:
            step += 1
            progress(0.7 * step / total_steps, desc="Repairing skin and faces...")
        enhanced = enhanced_skin_repair(enhanced)
        
        # Also apply face restoration specifically
        enhanced = restore_faces(enhanced)
    
    # Upscale for all modes except Face Enhance (which already upscales)
    if mode in (ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
        if progress:
            step += 1
            progress(0.8 * step / total_steps, desc="Upscaling image...")
        enhanced = advanced_tile_upscale(enhanced)
    
    # Detail Enhancement
    if mode in (ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
        if progress:
            step += 1
            progress(0.9 * step / total_steps, desc="Enhancing details...")
        enhanced = add_ultra_detail(enhanced, strength=0.7)
        enhanced = apply_high_frequency_details(enhanced, amount=0.5)
        enhanced = smart_sharpen(enhanced, strength=SMART_SHARPENING_STRENGTH)
    
        if progress:
            step += 1
            progress(0.95 * step / total_steps, desc="Adding final grain...")
        enhanced = add_film_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]
    # Gallery items can be filepath strings or (filepath, label) tuples.
    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, "
                  f"Reserved={reserved/1024**3:.2f}GB, "
                  f"Allocated={allocated/1024**3:.2f}GB, "
                  f"Free={free/1024**3:.2f}GB")
            
            return free > 1024**3  # Return True if more than 1GB free
        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
    """
    # Hardcode the negative prompt as requested
    negative_prompt = " "

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

    # Set up the generator for reproducibility
    generator = torch.Generator(device=device).manual_seed(seed)

    # Load input images into PIL Images — two optional slots.
    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

    # Fix for default 256x256 size
    if height == 256 and width == 256:
        height, width = None, None
    
    # Log generation parameters
    print(f"🎯 Generation Parameters:")
    print(f"   Prompt: '{prompt}'")
    print(f"   Negative Prompt: '{negative_prompt}'")
    print(f"   Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
    print(f"   Size: {width}x{height}, Images: {num_images_per_prompt}")
    print(f"   Enhance Mode: {enhance_mode}")

    # Check GPU memory before generation
    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.")

    # Generate the image
    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}")

    # Apply enhancement if requested
    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)
        ]

    # Save images to temporary files for proper serving
    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 after generation
    clear_gpu_cache()

    # Return image paths, seed, and make buttons visible when their feature is configured.
    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 Real-ESRGAN & Face Restoration</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
        - [Real-ESRGAN](https://huggingface.co/ai-forever/Real-ESRGAN) for high-quality upscaling
        - [GFPGAN](https://github.com/TencentARC/GFPGAN) for face restoration
        
        Upload an image and enter your prompt to edit it. The model uses your prompt exactly as provided.
        
        **💡 Pro Tips:**
        - Use **Face Enhance** mode for portrait photography
        - Use **Max Detail** for product shots and textures
        - Use **Full Enhance** for comprehensive improvement
        """)
        
        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"
                )
                
                # Enhancement info
                enhance_info = gr.Markdown("""
                **Enhancement Options:**
                - **Off**: No post-processing
                - **Upscale Only**: 4x upscaling with Real-ESRGAN
                - **Clean & Restore**: Artifact removal + skin/face restoration
                - **Max Detail**: Full detail enhancement with sharpening
                - **Face Enhance**: Specialized face restoration + upscaling
                - **Full Enhance**: Complete pipeline (clean + detail + face + upscale)
                """, visible=False, elem_classes="enhance-info")
                
                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
                    - Set custom dimensions for specific aspect ratios
                    """)

            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  # Made visible by default
                )

    # Event handlers
    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,
    )

    # Show enhancement info when enhance mode is changed
    enhance_mode.change(
        fn=lambda mode: gr.update(visible=mode != ENHANCE_MODE_OFF),
        inputs=[enhance_mode],
        outputs=[enhance_info]
    )

    # Use output as input button
    use_output_btn.click(
        fn=use_output_as_input,
        inputs=[result],
        outputs=[image_1]
    )

    # History gallery event handlers
    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__":
    # Check GPU availability
    print(f"🖥️  Device: {device}")
    if device == "cuda":
        print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
    
    # Check memory
    check_gpu_memory()
    
    # Launch the app
    print(f"🚀 Starting Pro Realism Edit Studio v{APP_VERSION}")
    print("=" * 60)
    demo.launch()