app.py

#!/usr/bin/env python3
"""
GPU-OPTIMIZED STYLE PAINTING APP - FIXED VERSION
Keeps ALL original features, just fixes what's broken

FEATURES PRESERVED:
✅ GPU acceleration with CUDA
✅ Multiple AI style models
✅ Real-time painting interface
✅ Preview vs AI processing distinction
✅ Auto-processing after delay
✅ Batch processing mode
✅ Pre-processed styles for speed
✅ NEW: Intensity control for each style

FIXES:
✅ Removed eraser (as requested) - just reset button
✅ Fixed Gradio update issues
✅ Each apply creates new base image
✅ Better state management
"""

import os
os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'

import gradio as gr
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
import cv2
import glob
import datetime
import tempfile
import time
import threading
import zipfile
import io
from typing import Dict, Tuple, Optional, List
import warnings
import traceback
warnings.filterwarnings("ignore")

# Force CUDA if available
if torch.cuda.is_available():
    torch.cuda.set_device(0)
    print("🔥 CUDA device set to 0")

# ===========================
# GPU SETUP (KEEP ORIGINAL)
# ===========================

def verify_gpu_setup():
    """Verify GPU is available and working"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("=" * 50)
    print("🔍 GPU VERIFICATION")
    print("=" * 50)
    print(f"🔍 CUDA Available: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"🔍 Device Name: {torch.cuda.get_device_name(0)}")
        total_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"🔍 Total GPU Memory: {total_memory:.2f} GB")
        print("✅ GPU Ready!")
    else:
        print("❌ CUDA NOT AVAILABLE - Running on CPU (will be slow)")
    
    print("=" * 50)
    return device

device = verify_gpu_setup()

# ===========================
# MODEL ARCHITECTURE (KEEP ORIGINAL)
# ===========================

class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        self.block = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features, affine=True),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features, affine=True)
        )

    def forward(self, x):
        return x + self.block(x)

class Generator(nn.Module):
    def __init__(self, input_nc=3, output_nc=3, n_residual_blocks=12):
        super(Generator, self).__init__()
        model = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(input_nc, 64, 7),
            nn.InstanceNorm2d(64, affine=True),
            nn.ReLU(inplace=True)
        ]
        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [
                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                nn.InstanceNorm2d(out_features, affine=True),
                nn.ReLU(inplace=True)
            ]
            in_features = out_features
            out_features = in_features * 2
        for _ in range(n_residual_blocks):
            model += [ResidualBlock(in_features)]
        out_features = in_features // 2
        for _ in range(2):
            model += [
                nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
                nn.InstanceNorm2d(out_features, affine=True),
                nn.ReLU(inplace=True)
            ]
            in_features = out_features
            out_features = in_features // 2
        model += [
            nn.ReflectionPad2d(3),
            nn.Conv2d(64, output_nc, 7),
            nn.Tanh()
        ]
        self.model = nn.Sequential(*model)

    def forward(self, x):
        return self.model(x)

class LegacyGenerator(nn.Module):
    def __init__(self, input_nc=3, output_nc=3):
        super(LegacyGenerator, self).__init__()
        model = [
            nn.ReflectionPad2d(3),
            nn.Conv2d(input_nc, 64, 7),
            nn.InstanceNorm2d(64, affine=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.InstanceNorm2d(128, affine=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.InstanceNorm2d(256, affine=True),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
            nn.InstanceNorm2d(128, affine=True),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.InstanceNorm2d(64, affine=True),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(3),
            nn.Conv2d(64, output_nc, 7),
            nn.Tanh()
        ]
        self.model = nn.Sequential(*model)

    def forward(self, x):
        return self.model(x)

def detect_model_architecture(state_dict):
    """Keep original model detection"""
    residual_keys = [key for key in state_dict.keys() if 'model.' in key and '.block.' in key]
    if not residual_keys:
        return 0
    block_indices = set()
    for key in residual_keys:
        try:
            parts = key.split('.')
            if len(parts) >= 3 and parts[2] == 'block':
                model_idx = int(parts[1])
                block_indices.add(model_idx)
        except (ValueError, IndexError):
            continue
    return len(block_indices) if block_indices else 12

def create_compatible_generator(state_dict):
    n_residual_blocks = detect_model_architecture(state_dict)
    if n_residual_blocks == 0:
        return LegacyGenerator()
    else:
        return Generator(n_residual_blocks=n_residual_blocks)

# ===========================
# FIXED PROCESSING SYSTEM
# ===========================

class StylePaintingSystem:
    def __init__(self):
        # Model management
        self.style_models = {}
        self.loaded_generators = {}
        self.precomputed_styles = {}
        
        # Current state - SIMPLIFIED WITHOUT ERASER
        self.original_image = None
        self.current_base = None  # Current working image (updates after each apply)
        self.current_display = None  # What's shown in UI
        self.is_preview = True
        
        # Painting state (NO ERASER)
        self.style_masks = {}  # style_key -> (mask, intensity)
        self.active_style = None
        self.active_intensity = 1.0  # NEW: intensity control
        
        # Processing state
        self.auto_timer = None
        self.auto_delay = 3.0
        self.processing_lock = threading.Lock()
        
        # Transforms (KEEP ORIGINAL)
        self.processing_size = 512
        self.transform = transforms.Compose([
            transforms.Resize((self.processing_size, self.processing_size)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        self.inverse_transform = transforms.Compose([
            transforms.Normalize((-1, -1, -1), (2, 2, 2)),
            transforms.ToPILImage()
        ])
        
        # Style configurations (KEEP ORIGINAL)
        self.style_configs = [
            ('🌸', 'Natural Bokeh', (255, 182, 193)),
            ('🌅', 'Natural Golden', (255, 215, 0)),
            ('☀️', 'Photo Golden', (255, 165, 0)),
            ('🌙', 'Day to Night', (25, 25, 112)),
            ('❄️', 'Summer Winter', (173, 216, 230)),
            ('📷', 'Photo Bokeh', (255, 228, 196)),
            ('🖼️', 'Photo Monet', (144, 238, 144)),
            ('🎨', 'Photo Seurat', (221, 160, 221)),
            ('🌫️', 'Foggy Clear', (220, 220, 220))
        ]
        
        self.discover_models()

    def discover_models(self):
        """Keep original model discovery"""
        print("\n" + "="*60)
        print("🔍 DISCOVERING MODELS")
        print("="*60)
        
        patterns = [
            './models/*_best_*/',
            './models/*/',
            './models/*.pth',
            './models/*/*.pth'
        ]
        
        all_files = []
        for pattern in patterns:
            files = glob.glob(pattern)
            if files:
                print(f"📂 Pattern '{pattern}' found {len(files)} items")
                all_files.extend(files)
        
        discovered = []
        
        for path in all_files:
            if os.path.isdir(path):
                ab_files = glob.glob(os.path.join(path, '*generator_AB.pth'))
                if ab_files:
                    folder_name = os.path.basename(path.rstrip('/'))
                    model_name = folder_name.split('_best_')[0] if '_best_' in folder_name else folder_name
                    discovered.append((model_name, ab_files[0]))
            elif path.endswith('.pth') and 'generator' in path:
                model_name = os.path.basename(path).replace('.pth', '').replace('_generator_AB', '')
                discovered.append((model_name, path))
        
        for idx, (model_name, model_path) in enumerate(discovered[:len(self.style_configs)]):
            emoji, display_name, color = self.style_configs[idx]
            
            self.style_models[model_name] = {
                'path': model_path,
                'emoji': emoji,
                'name': display_name,
                'color': color
            }
            print(f"📌 Registered: {emoji} {display_name} ({model_name})")
        
        print(f"\n✅ Registered {len(self.style_models)} style models")
        print("="*60 + "\n")

    def load_generator(self, model_key):
        """Keep original model loading"""
        if model_key in self.loaded_generators:
            return self.loaded_generators[model_key]
            
        if model_key not in self.style_models:
            return None
            
        try:
            model_path = self.style_models[model_key]['path']
            state_dict = torch.load(model_path, map_location=device)
            
            if 'generator' in state_dict:
                state_dict = state_dict['generator']
            
            generator = create_compatible_generator(state_dict)
            generator.load_state_dict(state_dict)
            generator.eval()
            generator = generator.to(device)
            
            if device.type == 'cuda':
                try:
                    generator = generator.half()
                except:
                    generator = generator.float()
            
            self.loaded_generators[model_key] = generator
            return generator
            
        except Exception as e:
            print(f"❌ Error loading {model_key}: {e}")
            return None

    def process_image_with_style(self, image, model_key, intensity=1.0):
        """Process with intensity control"""
        generator = self.load_generator(model_key)
        if generator is None:
            return None
            
        try:
            original_size = image.size
            img_resized = image.resize((self.processing_size, self.processing_size), Image.LANCZOS)
            img_tensor = self.transform(img_resized).unsqueeze(0).to(device)
            
            if device.type == 'cuda' and next(generator.parameters()).dtype == torch.float16:
                img_tensor = img_tensor.half()
            
            with torch.no_grad():
                if device.type == 'cuda':
                    torch.cuda.synchronize()
                    
                result_tensor = generator(img_tensor)
                
                if device.type == 'cuda':
                    torch.cuda.synchronize()
                
                if result_tensor.dtype == torch.float16:
                    result_tensor = result_tensor.float()
                
                result_tensor = result_tensor.cpu()
                processed_img = self.inverse_transform(result_tensor.squeeze(0))
            
            processed_img = processed_img.resize(original_size, Image.LANCZOS)
            
            # Apply intensity
            if intensity < 1.0:
                processed_array = np.array(processed_img, dtype=np.float32)
                original_array = np.array(image, dtype=np.float32)
                blended = original_array * (1 - intensity) + processed_array * intensity
                processed_img = Image.fromarray(blended.astype(np.uint8))
            
            return processed_img
            
        except Exception as e:
            print(f"❌ Error processing: {e}")
            return None

    def setup_new_image(self, image, progress_callback=None):
        """Setup without precomputing all styles"""
        if image is None:
            return "No image provided"
        
        if image.mode == 'RGBA':
            rgb_image = Image.new('RGB', image.size, (255, 255, 255))
            rgb_image.paste(image, mask=image.split()[3])
            image = rgb_image
        elif image.mode != 'RGB':
            image = image.convert('RGB')
        
        self.original_image = image
        self.current_base = image.copy()  # NEW: track current base
        self.current_display = image
        self.is_preview = True
        self.style_masks = {}
        self.precomputed_styles = {}
        
        # Optionally precompute SOME styles
        if progress_callback:
            progress_callback(0.5, "Image ready! Precomputing popular styles...")
            # Precompute just first 3 styles
            for idx, (model_key, model_data) in enumerate(list(self.style_models.items())[:3]):
                processed = self.process_image_with_style(image, model_key)
                if processed:
                    self.precomputed_styles[model_key] = processed
        
        return f"✅ Ready! {len(self.style_models)} styles available on {device.type.upper()}"

    def set_active_style(self, style_choice):
        """Set style (no eraser mode)"""
        for model_key, model_data in self.style_models.items():
            if f"{model_data['emoji']} {model_data['name']}" == style_choice:
                self.active_style = model_key
                return f"🎯 Selected: {style_choice}"
        
        return "❌ Style not found"

    def set_intensity(self, intensity):
        """Set painting intensity"""
        self.active_intensity = intensity
        return f"Intensity: {int(intensity * 100)}%"

    def get_painted_mask(self, editor_data):
        """Keep original mask extraction"""
        if not editor_data or not editor_data.get('layers'):
            return None
            
        if not self.current_base:
            return None
            
        height, width = self.current_base.size[1], self.current_base.size[0]
        combined_mask = np.zeros((height, width), dtype=bool)
        
        for layer in editor_data['layers']:
            if layer is None:
                continue
                
            layer_array = np.array(layer)
            
            if layer_array.shape[:2] != (height, width):
                layer_pil = Image.fromarray(layer_array)
                layer_pil = layer_pil.resize((width, height), Image.LANCZOS)
                layer_array = np.array(layer_pil)
            
            if layer_array.shape[-1] == 4:
                mask = layer_array[:, :, 3] > 30
            else:
                mask = np.any(layer_array > 10, axis=2)
            
            combined_mask = combined_mask | mask
        
        return combined_mask

    def update_painting(self, editor_data):
        """Update painting (no eraser logic)"""
        if not self.current_base:
            return self.current_base, "Upload an image first"
        
        painted_mask = self.get_painted_mask(editor_data)
        if painted_mask is None:
            return self.current_display, "No painting detected"
        
        if not self.active_style:
            return self.current_display, "Select a style first"
        
        # Store mask with intensity
        if self.active_style not in self.style_masks:
            self.style_masks[self.active_style] = (np.zeros_like(painted_mask), self.active_intensity)
        
        # Update mask
        existing_mask, _ = self.style_masks[self.active_style]
        new_pixels = painted_mask & (~existing_mask)
        
        if np.any(new_pixels):
            updated_mask = existing_mask | new_pixels
            self.style_masks[self.active_style] = (updated_mask, self.active_intensity)
            
            self.is_preview = True
            self.schedule_auto_process()
            preview = self.create_preview()
            
            style_name = self.style_models[self.active_style]['name']
            return preview, f"🎨 Added {np.sum(new_pixels):,} pixels of {style_name} at {int(self.active_intensity*100)}%"
        
        return self.current_display, "Paint in new areas"

    def create_preview(self):
        """Create preview from current base"""
        if not self.current_base:
            return None
            
        result = np.array(self.current_base, dtype=np.float32)
        
        for style_key, (mask, intensity) in self.style_masks.items():
            if not np.any(mask):
                continue
                
            if style_key not in self.style_models:
                continue
                
            color = np.array(self.style_models[style_key]['color'], dtype=np.float32)
            mask_smooth = cv2.GaussianBlur(mask.astype(np.float32), (15, 15), 0)
            mask_3d = np.stack([mask_smooth] * 3, axis=2)
            
            overlay = result * 0.7 + color * 0.3
            result = result * (1 - mask_3d * intensity) + overlay * mask_3d * intensity
        
        self.current_display = Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))
        return self.current_display

    def apply_ai_processing(self, progress_callback=None):
        """Apply and create new base - FIXED FOR GRADIO UPDATE"""
        if not self.current_base:
            return self.current_base, "No image loaded"
        
        has_styles = any(np.any(mask) for mask, _ in self.style_masks.values())
        if not has_styles:
            return self.current_base, "Nothing to process - paint first!"
        
        with self.processing_lock:
            # Start from current base (not original!)
            result = np.array(self.current_base, dtype=np.float32)
            applied_count = 0
            
            for style_key, (mask, intensity) in self.style_masks.items():
                if not np.any(mask):
                    continue
                
                # Process or use cached
                if style_key in self.precomputed_styles:
                    styled_img = self.precomputed_styles[style_key]
                else:
                    styled_img = self.process_image_with_style(self.current_base, style_key, intensity)
                    if styled_img:
                        self.precomputed_styles[style_key] = styled_img
                
                if styled_img:
                    styled_array = np.array(styled_img, dtype=np.float32)
                    mask_smooth = cv2.GaussianBlur(mask.astype(np.float32), (21, 21), 0)
                    mask_3d = np.stack([mask_smooth] * 3, axis=2)
                    
                    # Apply with intensity
                    result = result * (1 - mask_3d * intensity) + styled_array * mask_3d * intensity
                    applied_count += 1
            
            # Create new base image - THIS IS THE KEY FIX
            new_image = Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))
            
            # Update current base for next round
            self.current_base = new_image.copy()
            self.current_display = new_image
            self.is_preview = False
            
            # Clear masks for next painting session
            self.style_masks = {}
            self.precomputed_styles = {}  # Clear cache since base changed
            
            # Cancel auto-timer
            if self.auto_timer:
                self.auto_timer.cancel()
                self.auto_timer = None
            
            # Force new object for Gradio
            return new_image, f"🔥 Applied {applied_count} styles! Ready for more painting."

    def schedule_auto_process(self):
        """Keep original auto-processing"""
        if self.auto_timer:
            self.auto_timer.cancel()
        
        self.auto_timer = threading.Timer(self.auto_delay, self.auto_process)
        self.auto_timer.daemon = True
        self.auto_timer.start()

    def auto_process(self):
        """Auto process callback"""
        self.apply_ai_processing()

    def reset_all(self):
        """Reset to original image"""
        if self.original_image:
            self.current_base = self.original_image.copy()
            self.current_display = self.original_image
            self.style_masks = {}
            self.precomputed_styles = {}
            self.is_preview = True
            return self.original_image, "🔄 Reset to original"
        return None, "No image loaded"

    def process_batch(self, images, selected_styles_with_intensity, progress_callback=None):
        """Batch process with intensity"""
        results = {}
        total = len(images) * len(selected_styles_with_intensity)
        current = 0
        
        for img_idx, image in enumerate(images):
            img_results = {}
            
            for style_text, intensity in selected_styles_with_intensity:
                current += 1
                if progress_callback:
                    progress_callback(current / total, f"Processing image {img_idx+1} with {style_text} at {int(intensity*100)}%")
                
                # Find the model key
                model_key = None
                for key, data in self.style_models.items():
                    if f"{data['emoji']} {data['name']}" == style_text:
                        model_key = key
                        break
                
                if model_key:
                    processed = self.process_image_with_style(image, model_key, intensity)
                    if processed:
                        img_results[f"{style_text}_{int(intensity*100)}"] = processed
            
            results[f"image_{img_idx}"] = img_results
        
        return results

# ===========================
# GLOBAL SYSTEM INSTANCE
# ===========================

system = StylePaintingSystem()

# ===========================
# GRADIO INTERFACE FUNCTIONS (FIXED)
# ===========================

def on_image_upload(image, progress=gr.Progress()):
    """Handle image upload"""
    if image is None:
        return None, "Please upload an image", gr.update(value=None)
    
    def progress_cb(val, desc):
        progress(val, desc)
    
    status = system.setup_new_image(image, progress_cb)
    
    # Return image to both display and editor
    return image, status, image

def on_style_select(style_choice):
    """Handle style selection"""
    if not style_choice:
        return "Select a style"
    
    return system.set_active_style(style_choice)

def on_intensity_change(intensity):
    """Handle intensity change"""
    return system.set_intensity(intensity)

def on_paint_change(editor_data):
    """Handle painting changes"""
    if not editor_data:
        return system.current_display, "Paint to add styles"
    
    result_img, status = system.update_painting(editor_data)
    return result_img, status

def on_apply_ai():
    """Apply AI processing - FIXED FOR IMMEDIATE UPDATE"""
    # Cancel any auto-processing
    if system.auto_timer:
        system.auto_timer.cancel()
        system.auto_timer = None
    
    # Apply processing
    result_img, status = system.apply_ai_processing()
    
    # Force new object for Gradio
    if result_img:
        # Create a completely new image object
        import io
        buffer = io.BytesIO()
        result_img.save(buffer, format='PNG')
        buffer.seek(0)
        new_img = Image.open(buffer)
        
        # Return new image to BOTH displays to sync them
        return new_img, status, new_img
    
    return None, "Processing failed", None

def on_reset():
    """Reset to original"""
    result_img, status = system.reset_all()
    return result_img, status, result_img

def on_batch_process(file_list, selected_styles, intensities, progress=gr.Progress()):
    """Batch processing with intensity"""
    if not file_list or not selected_styles:
        return None, [], "Upload images and select styles"
    
    # Pair styles with intensities
    styles_with_intensity = []
    for i, style in enumerate(selected_styles):
        intensity = intensities[i] if i < len(intensities) else 1.0
        styles_with_intensity.append((style, intensity))
    
    # Load images
    images = []
    for file_path in file_list:
        try:
            img = Image.open(file_path)
            if img.mode == 'RGBA':
                rgb_img = Image.new('RGB', img.size, (255, 255, 255))
                rgb_img.paste(img, mask=img.split()[3])
                img = rgb_img
            images.append(img)
        except Exception as e:
            print(f"Error loading {file_path}: {e}")
    
    if not images:
        return None, [], "Failed to load images"
    
    def progress_cb(val, desc):
        progress(val, desc)
    
    # Process
    results = system.process_batch(images, styles_with_intensity, progress_cb)
    
    # Create outputs
    preview_images = []
    zip_buffer = io.BytesIO()
    
    with zipfile.ZipFile(zip_buffer, 'w') as zf:
        for img_key, img_results in results.items():
            for style_name, processed_img in img_results.items():
                filename = f"{img_key}_{style_name.replace(' ', '_')}.png"
                preview_images.append(processed_img)
                
                img_buffer = io.BytesIO()
                processed_img.save(img_buffer, format='PNG')
                zf.writestr(filename, img_buffer.getvalue())
    
    # Save zip
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    zip_path = os.path.join(tempfile.gettempdir(), f"batch_{timestamp}.zip")
    
    with open(zip_path, 'wb') as f:
        f.write(zip_buffer.getvalue())
    
    return zip_path, preview_images[:8], f"✅ Processed {len(images)} images"

# ===========================
# CREATE GRADIO INTERFACE (FIXED)
# ===========================

def create_interface():
    """Create the Gradio interface with fixes"""
    
    with gr.Blocks(title="GPU Style Painting", css="""
        .paint-canvas { height: 600px !important; }
    """) as interface:
        
        gr.Markdown("# 🎨 GPU-Optimized Style Painting")
        gr.Markdown(f"**Device:** {device.type.upper()} | **Auto-process:** {system.auto_delay}s")
        
        with gr.Tabs():
            # PAINTING TAB
            with gr.TabItem("🎨 Paint Mode"):
                with gr.Row():
                    with gr.Column(scale=1):
                        input_image = gr.Image(
                            type="pil",
                            label="Upload Image"
                        )
                        
                        upload_status = gr.Textbox(
                            label="Status",
                            value="Upload an image to start"
                        )
                        
                        # Style selection (NO ERASER)
                        style_choices = []
                        for key, data in system.style_models.items():
                            style_choices.append(f"{data['emoji']} {data['name']}")
                        
                        style_selector = gr.Radio(
                            choices=style_choices,
                            label="Select Style",
                            value=style_choices[0] if style_choices else None
                        )
                        
                        style_status = gr.Textbox(
                            label="Style Status",
                            value="Select a style to paint"
                        )
                        
                        # NEW: Intensity slider
                        intensity_slider = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=1.0,
                            step=0.1,
                            label="Style Intensity"
                        )
                        
                        intensity_status = gr.Textbox(
                            label="Intensity",
                            value="Intensity: 100%"
                        )
                        
                        # Control buttons
                        with gr.Row():
                            apply_btn = gr.Button("🔥 Apply AI", variant="primary")
                            reset_btn = gr.Button("🔄 Reset", variant="secondary")
                    
                    with gr.Column(scale=2):
                        # FIXED: Separate display and editor
                        result_image = gr.Image(
                            label="Result",
                            height=600,
                            type="pil",
                            interactive=False  # Display only
                        )
                        
                        painting_canvas = gr.ImageEditor(
                            type="pil",
                            label="Paint Canvas",
                            brush=gr.Brush(default_size=30),
                            height=600,
                            image_mode="RGB"
                        )
                        
                        paint_status = gr.Textbox(
                            label="Paint Status",
                            value="Ready"
                        )
            
            # BATCH TAB (WITH INTENSITY)
            with gr.TabItem("📦 Batch Mode"):
                with gr.Row():
                    with gr.Column():
                        batch_files = gr.File(
                            label="Upload Images",
                            file_count="multiple",
                            file_types=["image"]
                        )
                        
                        batch_styles = gr.CheckboxGroup(
                            choices=[f"{data['emoji']} {data['name']}" for data in system.style_models.values()],
                            label="Select Styles"
                        )
                        
                        # NEW: Intensity for each style
                        batch_intensities = gr.Dataframe(
                            headers=["Style", "Intensity"],
                            datatype=["str", "number"],
                            col_count=(2, "fixed"),
                            value=[["Style 1", 1.0], ["Style 2", 1.0], ["Style 3", 1.0]],
                            label="Style Intensities (0.1-1.0)"
                        )
                        
                        batch_btn = gr.Button("🚀 Process Batch", variant="primary")
                        batch_status = gr.Textbox(label="Status")
                    
                    with gr.Column():
                        batch_download = gr.File(label="Download ZIP")
                        batch_preview = gr.Gallery(
                            label="Preview",
                            columns=4,
                            height=400
                        )
        
        # Wire up events
        input_image.change(
            fn=on_image_upload,
            inputs=[input_image],
            outputs=[result_image, upload_status, painting_canvas]
        )
        
        style_selector.change(
            fn=on_style_select,
            inputs=[style_selector],
            outputs=[style_status]
        )
        
        intensity_slider.change(
            fn=on_intensity_change,
            inputs=[intensity_slider],
            outputs=[intensity_status]
        )
        
        painting_canvas.change(
            fn=on_paint_change,
            inputs=[painting_canvas],
            outputs=[result_image, paint_status],
            queue=False  # Immediate preview
        )
        
        apply_btn.click(
            fn=on_apply_ai,
            inputs=[],
            outputs=[result_image, paint_status, painting_canvas],
            queue=False  # IMMEDIATE execution
        )
        
        reset_btn.click(
            fn=on_reset,
            inputs=[],
            outputs=[result_image, paint_status, painting_canvas]
        )
        
        # Batch processing
        def prepare_batch_intensities(styles):
            """Create intensity dataframe based on selected styles"""
            return [[style, 1.0] for style in styles]
        
        batch_styles.change(
            fn=prepare_batch_intensities,
            inputs=[batch_styles],
            outputs=[batch_intensities]
        )
        
        def process_batch_with_df(files, styles, intensity_df):
            """Process batch using dataframe intensities"""
            intensities = [row[1] for row in intensity_df.values if row[0] in styles]
            return on_batch_process(files, styles, intensities)
        
        batch_btn.click(
            fn=process_batch_with_df,
            inputs=[batch_files, batch_styles, batch_intensities],
            outputs=[batch_download, batch_preview, batch_status]
        )
        
        gr.Markdown("""
        ## 📖 Instructions
        1. **Upload** an image to start
        2. **Select** a style and adjust intensity
        3. **Paint** on the canvas - see instant preview
        4. **Apply AI** to process the painted areas
        5. **Continue** painting more styles - they build on previous results
        6. **Reset** to return to original image
        
        **Features:**
        - 🚀 GPU accelerated processing
        - 🎨 Multiple AI styles with intensity control
        - ⚡ Instant preview with color overlays
        - 🔥 Progressive application (each apply builds on previous)
        - 📦 Batch processing with per-style intensity
        """)
        
        return interface

# ===========================
# LAUNCH
# ===========================

if __name__ == "__main__":
    print("🚀 Starting GPU Style Painting App...")
    print(f"📍 Device: {device}")
    print(f"⏰ Auto-process delay: {system.auto_delay}s")
    
    interface = create_interface()
    interface.queue()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )