#!/usr/bin/env python3
"""
Gradio Demo App for Patchioner Model - Trace-based Image Captioning

This demo allows users to:
1. Upload or select an image
2. Draw traces on the image using Gradio's ImageEditor
3. Generate captions for the traced regions using a pre-trained Patchioner model

Author: Generated for decap-dino project
"""

import os

import gradio as gr

from gradio_image_annotation import image_annotator as foo_image_annotator

import torch
import yaml
import traceback
from pathlib import Path
from PIL import Image
import numpy as np
from typing import List, Dict

# Import the Patchioner model from the src directory
from src.model import Patchioner

# Global variable to store the loaded model
loaded_model = None
model_config_path = None
device = "cuda" if torch.cuda.is_available() else "cpu"

# Default model configuration
DEFAULT_MODEL_CONFIG = "mlp.viecap.k.yaml"

# Example images directory
current_dir = os.path.dirname(__file__)
EXAMPLE_IMAGES_DIR = Path(os.path.join(current_dir, 'example-images')).resolve()
CONFIGS_DIR = Path(os.path.join(current_dir, 'configs')).resolve()


def initialize_default_model() -> str:
    """Initialize the default model at startup."""
    global loaded_model, model_config_path
    
    try:
        # Look for the default config file
        default_config_path = CONFIGS_DIR / DEFAULT_MODEL_CONFIG
        
        if not default_config_path.exists():
            return f"❌ Default config file not found: {default_config_path}"
        
        print(f"Loading default model: {DEFAULT_MODEL_CONFIG}")
        
        # Load and parse the config
        with open(default_config_path, 'r') as f:
            config = yaml.safe_load(f)
        
        # Load the model using the from_config class method
        model = Patchioner.from_config(config, device=device)
        model.eval()
        model.to(device)
        
        # Store the model globally
        loaded_model = model
        model_config_path = str(default_config_path)
        
        return f"✅ Default model loaded: {DEFAULT_MODEL_CONFIG} on {device}"
        
    except Exception as e:
        error_msg = f"❌ Error loading default model: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return error_msg


def get_example_images(limit=None) -> List[str]:
    """Get list of example images for the demo."""
    example_images = []
    if EXAMPLE_IMAGES_DIR.exists():
        for ext in ['*.jpg', '*.jpeg', '*.png']:
            example_images.extend(str(p) for p in EXAMPLE_IMAGES_DIR.glob(ext))
    if limit is not None:
        example_images = example_images[:limit]
    return example_images
    

def get_example_configs() -> List[str]:
    """Get list of example config files."""
    example_configs = []
    if CONFIGS_DIR.exists():
        example_configs = [str(p) for p in CONFIGS_DIR.glob("*.yaml")]
    else:
        print(f"Warning: Configs directory {CONFIGS_DIR} does not exist.")
    return sorted(example_configs)


def load_model_from_config(config_file_path: str) -> str:
    """
    Load the Patchioner model from a config file.
    
    Args:
        config_file_path: Path to the YAML configuration file
        
    Returns:
        Status message about model loading
    """
    global loaded_model, model_config_path
    
    try:
        if not config_file_path or not os.path.exists(config_file_path):
            return "❌ Error: Config file path is empty or file does not exist."
        
        print(f"Loading model from config: {config_file_path}")
        
        # Load and parse the config
        with open(config_file_path, 'r') as f:
            config = yaml.safe_load(f)
        
        # Load the model using the from_config class method
        model = Patchioner.from_config(config, device=device)
        model.eval()
        model.to(device)
        
        # Store the model globally
        loaded_model = model
        model_config_path = config_file_path
        
        return f"✅ Model loaded successfully from {os.path.basename(config_file_path)} on {device}"
        
    except Exception as e:
        error_msg = f"❌ Error loading model: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return error_msg


def process_image_trace_to_coordinates(image_editor_data) -> List[List[Dict[str, float]]]:
    """
    Convert Gradio ImageEditor trace data to the coordinate format expected by the model.
    
    The expected format is: [[{"x": float, "y": float, "t": float}, ...], ...]
    where coordinates are normalized to [0, 1] and t is a timestamp.
    
    Args:
        image_editor_data: Data from Gradio ImageEditor component
        
    Returns:
        List of traces in the expected format
    """
    try:
        print(f"[DEBUG] process_image_trace_to_coordinates called")
        print(f"[DEBUG] image_editor_data type: {type(image_editor_data)}")
        
        if image_editor_data is None:
            print("[DEBUG] image_editor_data is None")
            return []
            
        if isinstance(image_editor_data, dict):
            print(f"[DEBUG] Available keys in image_editor_data: {list(image_editor_data.keys())}")
        
        # Check for different possible structures
        layers = None
        if isinstance(image_editor_data, dict):
            if 'layers' in image_editor_data:
                layers = image_editor_data['layers']
            elif 'composite' in image_editor_data:
                # Sometimes gradio stores drawing data differently
                composite = image_editor_data['composite']
                if isinstance(composite, dict) and 'layers' in composite:
                    layers = composite['layers']
        
        if not layers:
            print("[DEBUG] No layers found in image_editor_data")
            return []
        
        traces = []
        print(f"[DEBUG] Processing {len(layers)} layers")
        
        # Process each drawing layer - they are PIL Images, not coordinate data
        for i, layer in enumerate(layers):
            print(f"[DEBUG] Processing layer {i}: {layer}")
            
            # Skip if layer is not a PIL Image or is empty
            if not isinstance(layer, Image.Image):
                print(f"[DEBUG] Layer {i} is not a PIL Image")
                continue
            
            # Convert layer to numpy array to find non-transparent pixels
            layer_array = np.array(layer)
            
            # Find non-transparent pixels (alpha > 0)
            if layer_array.shape[2] == 4:  # RGBA
                non_transparent = layer_array[:, :, 3] > 0
            else:  # RGB - assume any non-black pixel is drawn
                non_transparent = np.any(layer_array > 0, axis=2)
            
            # Get coordinates of drawn pixels
            y_coords, x_coords = np.where(non_transparent)
            
            if len(x_coords) == 0:
                print(f"[DEBUG] Layer {i} has no drawn pixels")
                continue
            
            print(f"[DEBUG] Layer {i} has {len(x_coords)} drawn pixels")
            
            # Convert pixel coordinates to trace format
            trace_points = []
            img_height, img_width = layer_array.shape[:2]
            
            # Sample some points from the drawn pixels (to avoid too many points)
            num_points = min(len(x_coords), 100)  # Limit to 100 points max
            if num_points > 0:
                # Sample evenly spaced indices
                indices = np.linspace(0, len(x_coords) - 1, num_points, dtype=int)
                sampled_x = x_coords[indices]
                sampled_y = y_coords[indices]
                
                # Convert to normalized coordinates and create trace points
                for idx, (x, y) in enumerate(zip(sampled_x, sampled_y)):
                    # Normalize coordinates to [0, 1]
                    x_norm = float(x) / img_width if img_width > 0 else 0
                    y_norm = float(y) / img_height if img_height > 0 else 0
                    
                    # Clamp to [0, 1] range
                    x_norm = max(0, min(1, x_norm))
                    y_norm = max(0, min(1, y_norm))
                    
                    # Add timestamp (arbitrary progression)
                    t = idx * 0.1
                    
                    trace_points.append({
                        "x": x_norm,
                        "y": y_norm,
                        "t": t
                    })
            
            if trace_points:
                traces.append(trace_points)
        
        return traces
        
    except Exception as e:
        print(f"Error processing image trace: {e}")
        print(traceback.format_exc())
        return []


def process_bounding_box_coordinates(annotator_data) -> List[List[float]]:
    """
    Convert Gradio image_annotator data to bounding box format expected by the model.
    
    Args:
        annotator_data: Data from Gradio image_annotator component
        
    Returns:
        List of bounding boxes in [x, y, width, height] format
    """
    try:
        print(f"[DEBUG] process_bounding_box_coordinates called")
        print(f"[DEBUG] annotator_data type: {type(annotator_data)}")
        #print(f"[DEBUG] annotator_data content: {annotator_data}")
        
        if annotator_data is None:
            print("[DEBUG] annotator_data is None")
            return []
            
        boxes = []
        
        # Handle the dictionary format from image_annotator
        if isinstance(annotator_data, dict):
            print(f"[DEBUG] Available keys in annotator_data: {list(annotator_data.keys())}")
            
            # Extract boxes from the 'boxes' key
            if 'boxes' in annotator_data and annotator_data['boxes']:
                for box in annotator_data['boxes']:
                    if isinstance(box, dict):
                        # Based on image_annotator.py, boxes have format:
                        # {"xmin": x, "ymin": y, "xmax": x2, "ymax": y2, "label": ..., "color": ...}
                        xmin = box.get('xmin', 0)
                        ymin = box.get('ymin', 0)
                        xmax = box.get('xmax', 0)
                        ymax = box.get('ymax', 0)
                        
                        width = xmax - xmin
                        height = ymax - ymin
                        
                        # Convert to [x, y, width, height] format
                        boxes.append([xmin, ymin, width, height])
            else:
                print("[DEBUG] No 'boxes' key found or boxes list is empty")
        
        print(f"[DEBUG] Found {len(boxes)} bounding boxes: {boxes}")
        return boxes
        
    except Exception as e:
        print(f"Error processing bounding box: {e}")
        print(traceback.format_exc())
        return []


def generate_caption(mode, image_data) -> str:
    """
    Generate caption for the image and traces/bboxes using the loaded model.
    
    Args:
        mode: Either "trace" or "bbox" mode
        image_data: Data from Gradio ImageEditor or Annotate component
        
    Returns:
        Generated caption or error message
    """
    global loaded_model
    
    try:
        print(f"[DEBUG] generate_caption called with mode: {mode}")
        print(f"[DEBUG] image_data type: {type(image_data)}")
        print(f"[DEBUG] image_data content: {image_data}")
        
        if loaded_model is None:
            return "❌ Error: No model loaded. Please load a model first using the config file."
        
        # Handle different input formats from Gradio components
        image = None
        if image_data is None:
            return "❌ Error: No image data provided."
        
        # Check if it's a PIL Image directly
        if isinstance(image_data, Image.Image):
            print("[DEBUG] Received PIL Image directly")
            image = image_data
        # Check if it's a dict (from image_annotator component)
        elif isinstance(image_data, dict):
            print(f"[DEBUG] Received dict with keys: {list(image_data.keys())}")
            if 'image' in image_data:
                image_array = image_data['image']
                # Convert numpy array to PIL Image if needed
                if hasattr(image_array, 'shape') and len(image_array.shape) == 3:
                    print("[DEBUG] Converting numpy array to PIL Image")
                    image = Image.fromarray(image_array)
                else:
                    image = image_array
            elif 'background' in image_data:
                image_array = image_data['background']
                # Convert numpy array to PIL Image if needed
                if hasattr(image_array, 'shape') and len(image_array.shape) == 3:
                    print("[DEBUG] Converting numpy array to PIL Image")
                    image = Image.fromarray(image_array)
                else:
                    image = image_array
            else:
                return f"❌ Error: No image found in data. Available keys: {list(image_data.keys())}"
        # Check for tuple/list format (from ImageEditor component)
        elif isinstance(image_data, (tuple, list)) and len(image_data) >= 1:
            print(f"[DEBUG] Received tuple/list with {len(image_data)} elements")
            image = image_data[0]  # First element should be the image
            if not isinstance(image, Image.Image):
                # Sometimes the structure might be different, search for PIL Image
                for item in image_data:
                    if isinstance(item, Image.Image):
                        image = item
                        break
        else:
            return f"❌ Error: Unexpected data type: {type(image_data)}"
        
        if image is None:
            return "❌ Error: Image is None."
        
        # Convert PIL image if necessary
        if not isinstance(image, Image.Image):
            return "❌ Error: Invalid image format."
        
        # Convert image to RGB if needed
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        if mode == "trace":
            return generate_trace_caption(image_data, image)
        elif mode == "bbox":
            return generate_bbox_caption(image_data, image)
        else:
            return f"❌ Error: Unknown mode: {mode}"
            
    except Exception as e:
        error_msg = f"❌ Error generating caption: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return error_msg


def generate_trace_caption(image_data, image) -> str:
    """Generate caption using traces."""
    global loaded_model
    
    try:
        # Process traces
        print("[DEBUG] Processing traces...")
        traces = process_image_trace_to_coordinates(image_data)
        print(f"[DEBUG] Found {len(traces)} traces")
        
        if not traces:
            # For debugging, let's generate a simple image caption instead of failing
            print("[DEBUG] No traces found, generating image caption instead")
            image_tensor = loaded_model.image_transforms(image).unsqueeze(0).to(device)
            
            with torch.no_grad():
                outputs = loaded_model(
                    image_tensor,
                    get_cls_capt=True,  # Get class caption as fallback
                    get_patch_capts=False,
                    get_avg_patch_capt=False
                )
            
            if 'cls_capt' in outputs:
                return f"🔍 No traces drawn. Image caption: {outputs['cls_capt']}"
            else:
                return "❌ Error: No traces detected. Please draw some traces on the image."
        
        print(f"Processing {len(traces)} traces")
        
        # Prepare image tensor
        image_tensor = loaded_model.image_transforms(image).unsqueeze(0).to(device)
        
        # Generate caption using the model
        with torch.no_grad():
            outputs = loaded_model(
                image_tensor,
                traces=traces,
                get_cls_capt=False,  # We want trace captions, not class captions
                get_patch_capts=False,
                get_avg_patch_capt=False
            )
        
        # Extract the trace captions
        if 'trace_capts' in outputs:
            captions = outputs['trace_capts']
            if isinstance(captions, list) and captions:
                captions = [cap.replace("<|startoftext|>", "").replace("<|endoftext|>", "") for cap in captions]
                # Join multiple captions if there are multiple traces
                if len(captions) == 1:
                    return f"Generated Caption: {captions[0]}"
                else:
                    formatted_captions = []
                    for i, caption in enumerate(captions, 1):
                        formatted_captions.append(f"Trace {i}: {caption}")
                    return "Generated Captions:\n" + "\n".join(formatted_captions)
            elif isinstance(captions, str):
                return f"Generated Caption: {captions}"
            else:
                return "❌ Error: No captions generated."
        else:
            return "❌ Error: Model did not return trace captions."
            
    except Exception as e:
        error_msg = f"❌ Error generating trace caption: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return error_msg


def generate_bbox_caption(image_data, image) -> str:
    """Generate caption using bounding boxes."""
    global loaded_model
    
    try:
        # Process bounding boxes
        print("[DEBUG] Processing bounding boxes...")
        bboxes = process_bounding_box_coordinates(image_data)
        print(f"[DEBUG] Found {len(bboxes)} bounding boxes")
        
        if not bboxes:
            # For debugging, let's generate a simple image caption instead of failing
            print("[DEBUG] No bounding boxes found, generating image caption instead")
            image_tensor = loaded_model.image_transforms(image).unsqueeze(0).to(device)
            
            with torch.no_grad():
                outputs = loaded_model(
                    image_tensor,
                    get_cls_capt=True,  # Get class caption as fallback
                    get_patch_capts=False,
                    get_avg_patch_capt=False
                )
            
            if 'cls_capt' in outputs:
                return f"🔍 No bounding boxes drawn. Image caption: {outputs['cls_capt']}"
            else:
                return "❌ Error: No bounding boxes detected. Please draw some bounding boxes on the image."
        
        print(f"Processing {len(bboxes)} bounding boxes")
        
        # Generate caption using the caption_bboxes method (as in eval_densecap.py)
        try:
            captions = loaded_model.caption_bboxes([image], [bboxes], crop_boxes=True)

            if isinstance(captions, list) and captions:
                if isinstance(captions[0], list):
                    captions = captions[0]  # Unwrap nested list if needed
                captions = [cap.replace("<|startoftext|>", "").replace("<|endoftext|>", "") for cap in captions]
                # Join multiple captions if there are multiple bboxes
                if len(captions) == 1:
                    return f"Generated Caption: {captions[0]}"
                else:
                    formatted_captions = []
                    for i, caption in enumerate(captions, 1):
                        formatted_captions.append(f"BBox {i}: {caption}")
                    return "Generated Captions:\n" + "\n".join(formatted_captions)
            elif isinstance(captions, str):
                return f"Generated Caption: {captions}"
            else:
                return "❌ Error: No captions generated."
                
        except Exception as e:
            print(f"Error using caption_bboxes method: {e}")
            # Fallback to regular forward method with bboxes
            image_tensor = loaded_model.image_transforms(image).unsqueeze(0).to(device)
            bbox_tensor = torch.tensor([bboxes]).to(device)
            
            with torch.no_grad():
                outputs = loaded_model(
                    image_tensor,
                    bboxes=bbox_tensor,
                    get_cls_capt=False,
                    get_patch_capts=False,
                    get_avg_patch_capt=False
                )
            
            if 'bbox_capts' in outputs:
                captions = outputs['bbox_capts']
                if isinstance(captions, list) and captions:
                    if isinstance(captions[0], list):
                        captions = captions[0]  # Unwrap nested list if needed
                    captions = [cap.replace("<|startoftext|>", "").replace("<|endoftext|>", "") for cap in captions]
                    if len(captions) == 1:
                        return f"Generated Caption: {captions[0]}"
                    else:
                        formatted_captions = []
                        for i, caption in enumerate(captions, 1):
                            formatted_captions.append(f"BBox {i}: {caption}")
                        return "Generated Captions:\n" + "\n".join(formatted_captions)
                elif isinstance(captions, str):
                    return f"Generated Caption: {captions}"
                else:
                    return "❌ Error: No captions generated."
            else:
                return "❌ Error: Model did not return bbox captions."
            
    except Exception as e:
        error_msg = f"❌ Error generating bbox caption: {str(e)}"
        print(error_msg)
        print(traceback.format_exc())
        return error_msg


def create_gradio_interface():
    """Create and configure the Gradio interface."""
    
    # Get example files
    example_images = get_example_images()
    example_configs = get_example_configs()


    custom_js = """
        <script>
        window.addEventListener("load", () => {
            // Hide Crop, Erase, and Color buttons
            const cropBtn   = document.querySelector('.image-editor__tool[title="Crop"]');
            const eraseBtn  = document.querySelector('.image-editor__tool[title="Erase"]');
            const colorBtn  = document.querySelector('.image-editor__tool[title="Color"]');

            [cropBtn, eraseBtn, colorBtn].forEach(btn => {
                console.log("Going to disable display for ", btn);
                if (btn) btn.style.display = "none";
            });

            // Optionally, select the Brush/Draft tool right away
            const brushBtn = document.querySelector('.image-editor__tool[title="Draw"]');
            console.log("Selecting brushbtn: ", brushBtn);
            if (brushBtn) brushBtn.click();
        });
        </script>
        """
    
    with gr.Blocks(
        title="Patchioner Trace Captioning Demo",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        """
    ) as demo:
        #gr.HTML(custom_js)  # inject custom JS
        
        gr.Markdown("""
        # 🎯 Patchioner Trace Captioning Demo
        
        This demo allows you to:
        1. **Select a captioning mode** (trace or bounding box)
        2. **Upload or select an image** from examples  
        3. **Draw traces or bounding boxes** on the image
        4. **Generate captions** describing the marked areas
        
        ## Instructions:
        1. Choose between Trace or BBox mode
        2. Upload an image or use one of the provided examples
        3. Use the appropriate tool to mark areas of interest in the image
        4. Click "Generate Caption" to get AI-generated descriptions
        
        **Model:** Using `mlp.karpathy.yaml` configuration (automatically loaded)
        """)
        
        # Initialize model status
        model_initialization_status = initialize_default_model()
        
        with gr.Row():
            gr.Markdown(f"**Model Status:** {model_initialization_status}")
        
        with gr.Row():
            mode_selector = gr.Radio(
                choices=["trace", "bbox"],
                value="trace",
                label="📋 Captioning Mode",
                info="Choose between trace-based or bounding box-based captioning",
                visible=False
            )
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 🖼️ Image Editor")
                
                # Image editor for drawing traces (default)
                image_editor = gr.ImageEditor(
                    label="Upload image and draw traces",
                    type="pil",
                    crop_size=None,
                    brush=gr.Brush(default_size=3, colors=["red", "blue", "green", "yellow", "purple"]),
                    visible=True,
                    #tools=["brush"],
                    height=600
                )
                
                # Image annotator for bounding boxes (hidden by default)
                image_annotator = foo_image_annotator( #gr.Image(
                    label="Upload image and draw bounding boxes",
                    visible=False,
                    #classes=["object"],
                    #type="bbox"
                    #tool="select"
                )
                
            with gr.Column():
                if example_images:
                    gr.Markdown("#### 📷 Or select from example images:")
                    example_gallery = gr.Gallery(
                        value=example_images,
                        label="Example Images",
                        show_label=True,
                        elem_id="gallery",
                        columns=3,
                        rows=2,
                        object_fit="contain",
                        height="auto"
                    )
        
        with gr.Row():
            generate_button = gr.Button("✨ Generate Caption", variant="primary", size="lg")
        
        with gr.Row():
            output_text = gr.Textbox(
                label="Generated Caption",
                placeholder="Generated caption will appear here...",
                lines=5,
                max_lines=10,
                interactive=False
            )
        
        # Event handlers
        def toggle_input_components(mode):
            """Toggle between image editor and annotator based on mode."""
            if mode == "trace":
                return gr.update(visible=True), gr.update(visible=False)
            else:  # bbox mode
                return gr.update(visible=False), gr.update(visible=True)
        
        def load_example_image_to_both(evt: gr.SelectData):
            """Load selected example image into both components."""
            try:
                example_images = get_example_images()
                if evt.index < len(example_images):
                    selected_image_path = example_images[evt.index]
                    img = Image.open(selected_image_path)
                    # For ImageEditor, return the PIL image directly
                    # For image_annotator, return dict format as expected by the component
                    annotated_format = {
                        "image": img,
                        "boxes": [],
                        "orientation": 0
                    }
                    return img, annotated_format
                return None, {"image": None, "boxes": [], "orientation": 0}
            except Exception as e:
                print(f"Error loading example image: {e}")
                return None, {"image": None, "boxes": [], "orientation": 0}
        
        def generate_caption_wrapper(mode, image_editor_data, image_annotator_data):
            """Wrapper to call generate_caption with the appropriate data based on mode."""
            if mode == "trace":
                return generate_caption(mode, image_editor_data)
            else:  # bbox mode
                return generate_caption(mode, image_annotator_data)
        
        # Connect event handlers
        mode_selector.change(
            fn=toggle_input_components,
            inputs=mode_selector,
            outputs=[image_editor, image_annotator]
        )
        
        generate_button.click(
            fn=generate_caption_wrapper,
            inputs=[mode_selector, image_editor, image_annotator],
            outputs=output_text
        )
        
        if example_images:
            example_gallery.select(
                fn=load_example_image_to_both,
                outputs=[image_editor, image_annotator]
            )
        
        gr.Markdown("""
        ### 💡 Tips:
        - **Mode Selection**: Switch between trace and bounding box modes based on your needs
        - **Trace Mode**: Draw continuous lines over areas you want to describe
        - **BBox Mode**: Draw rectangular bounding boxes around objects of interest
        - **Multiple Areas**: Create multiple traces/boxes for different objects to get individual captions
        - **Model Performance**: First load may take some time as weights are downloaded
        
        ### 🔧 Technical Details:
        - **Trace Mode**: Converts drawings to normalized (x, y) coordinates with timestamps
        - **BBox Mode**: Uses bounding box coordinates for region-specific captioning
        - **Model Architecture**: Uses `mlp.karpathy.yaml` configuration with CLIP and ViT components
        - **Processing**: Each trace/bbox is processed separately to generate corresponding captions
        """)
    
    return demo


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Patchioner Trace Captioning Demo")
    parser.add_argument("--port", type=int, default=4141, help="Port to run the Gradio app on")
    parser.add_argument("--local", action="store_true", help="Run the app locally. If not set, the app will be use default values for Gradio sharing host and ports.")
    args = parser.parse_args()

    print("Starting Patchioner Trace Captioning Demo...")
    print(f"Using device: {device}")
    print(f"Default model: {DEFAULT_MODEL_CONFIG}")
    print(f"Example images directory: {EXAMPLE_IMAGES_DIR}")
    print(f"Configs directory: {CONFIGS_DIR}")
    
    demo = create_gradio_interface()
    if not args.local:
        demo.launch()
    else:
        demo.launch(
            server_name="0.0.0.0",
            server_port=args.port,
            share=True,
            debug=True
        )