app.py

import base64
import json
import time
import os
import uuid
import threading
from io import BytesIO
import requests
import gradio as gr
from PIL import Image

# 全局变量存储应用实例
app_instance = None


def cleanup_temp_files(temp_dirs=None):
    """
    清理指定的临时文件夹中的旧文件
    temp_dirs: 要清理的目录列表，如果为None则只清理记录的目录
    """
    try:
        # 如果没有指定目录，只清理我们记录的目录
        if temp_dirs is None:
            temp_dirs = []
        
        current_time = time.time()
        cleaned_count = 0
        
        for temp_dir in temp_dirs:
            try:
                # 确保目录存在且是目录
                if not os.path.exists(temp_dir) or not os.path.isdir(temp_dir):
                    continue
                
                # 遍历目录中的文件
                for root, dirs, files in os.walk(temp_dir):
                    for file_name in files:
                        file_path = os.path.join(root, file_name)
                        try:
                            # 检查文件修改时间
                            file_mtime = os.path.getmtime(file_path)
                            # 如果文件超过30分钟未修改，则删除
                            if current_time - file_mtime > 1800:  # 1800秒 = 30分钟
                                # 检查是否是图片文件或临时文件
                                if any(file_path.lower().endswith(ext) for ext in ['.jpg', '.jpeg', '.png', '.gif', '.webp', '.tmp']):
                                    os.remove(file_path)
                                    cleaned_count += 1
                                    print(f"Cleaned old temp file: {file_path}")
                        except Exception as e:
                            print(f"Error cleaning {file_path}: {e}")
                    
                    # 清理空目录
                    for dir_name in dirs:
                        dir_path = os.path.join(root, dir_name)
                        try:
                            if os.path.exists(dir_path) and not os.listdir(dir_path):
                                os.rmdir(dir_path)
                                print(f"Removed empty temp directory: {dir_path}")
                        except Exception as e:
                            print(f"Error removing directory {dir_path}: {e}")
                            
            except Exception as e:
                print(f"Error processing directory {temp_dir}: {e}")
        
        if cleaned_count > 0:
            print(f"Cleanup completed: removed {cleaned_count} temporary files")
        
    except Exception as e:
        print(f"Error during temp cleanup: {e}")


# 全局变量记录使用过的临时目录
used_temp_dirs = set()


def start_cleanup_scheduler():
    """
    启动定时清理任务
    """
    def cleanup_worker():
        while True:
            try:
                # 每30分钟清理一次
                time.sleep(1800)  # 1800秒 = 30分钟
                print("Starting scheduled cleanup...")
                # 只清理我们记录的临时目录
                cleanup_temp_files(list(used_temp_dirs))
            except Exception as e:
                print(f"Error in cleanup scheduler: {e}")
    
    # 创建守护线程
    cleanup_thread = threading.Thread(target=cleanup_worker, daemon=True)
    cleanup_thread.start()


def upload_file_to_kie(file_path: str, api_key: str):
    """
    上传文件到 KIE AI 的文件存储服务
    返回文件的公开访问URL
    """
    url = "https://kieai.redpandaai.co/api/file-stream-upload"
    headers = {
        "Authorization": f"Bearer {api_key}"
    }

    # 获取文件名和扩展名
    file_name = os.path.basename(file_path)
    file_ext = os.path.splitext(file_name)[1].lower()
    if file_ext is None or file_ext == "":
        file_ext = ".jpg"
    file_name = uuid.uuid4().hex + file_ext

    # 根据文件扩展名确定MIME类型
    mime_types = {
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.gif': 'image/gif',
        '.webp': 'image/webp'
    }
    mime_type = mime_types.get(file_ext, 'image/jpg')

    try:
        with open(file_path, 'rb') as file_handle:
            # 准备文件上传数据
            files = {
                'file': (file_name, file_handle, mime_type)
            }
            data = {
                'uploadPath': 'images/nano_banana',
                "fileName": file_name
            }

            response = requests.post(url, headers=headers, files=files, data=data, timeout=60)
        if response.status_code == 200:
            resp_json = response.json()
            if resp_json.get("success") and resp_json.get("data"):
                # 返回下载URL
                return True, resp_json["data"]["downloadUrl"]
            else:
                return False, resp_json.get("msg", "Upload failed")
        else:
            return False, f"Upload error, please try again later"
    except:
        return False, f"Upload error, please try again later"


def create_task(prompt: str, image_urls: list[str], api_key: str):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    url = "https://api.kie.ai/api/v1/playground/createTask"
    payload = {
        "model": "google/nano-banana-edit",
        "callBackUrl": "",
        "input": {
            "prompt": prompt,
            "image_urls": image_urls
        }
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=60)
        if response.status_code == 200:
            resp_json = response.json()
            if resp_json.get("code") == 200:
                return 200, resp_json["data"]["taskId"]
            else:
                return 500, resp_json["message"]
    except Exception:
        pass
    return 500, "Internal Server Error"


def get_task_result(taskId: str, api_key: str):
    start_time = time.time()
    url = "https://api.kie.ai/api/v1/playground/recordInfo"
    params = {"taskId": taskId}
    headers = {"Authorization": f"Bearer {api_key}"}
    while time.time() - start_time < 600:
        try:
            response = requests.get(url, headers=headers, params=params, timeout=60)
            if response.status_code == 200:
                resp_json = response.json()
                if resp_json.get("code") == 200:
                    if resp_json['data']['state'] == 'success':
                        # 解析resultJson字符串
                        result_json = json.loads(resp_json['data']['resultJson'])
                        return result_json.get('resultUrls', [])  # 图片连接数组
                    elif resp_json['data']['state'] == 'fail':
                        return 500, resp_json['data']['failMsg']
                    else:
                        time.sleep(5)  # 等待2秒后再次查询
                        continue
                else:
                    return 500, resp_json["message"]
            else:
                break
        except Exception as e:
            pass
        time.sleep(2)  # 等待2秒后再次查询
    return 500, "Task timeout, please try again later"


def get_image_list(file_paths):
    """从文件路径列表获取PIL图片列表用于展示"""
    if not file_paths:
        return []

    images = []
    for path in file_paths:
        try:
            img = Image.open(path)
            images.append(img)
        except Exception:
            continue
    return images


def create_image_html(file_paths):
    """创建图片展示的HTML代码"""
    if not file_paths:
        return "<div id='image-display-area'><div class='no-images'>No images yet, please upload images</div></div>"

    html_items = []
    for i, path in enumerate(file_paths):
        try:
            # 将图片转换为base64编码
            with open(path, "rb") as img_file:
                img_data = base64.b64encode(img_file.read()).decode()
                img_ext = path.split('.')[-1].lower()
                if img_ext in ['jpg', 'jpeg']:
                    mime_type = 'image/jpeg'
                elif img_ext == 'png':
                    mime_type = 'image/png'
                elif img_ext == 'gif':
                    mime_type = 'image/gif'
                elif img_ext == 'webp':
                    mime_type = 'image/webp'
                else:
                    mime_type = 'image/jpeg'

                img_src = f"data:{mime_type};base64,{img_data}"

                html_items.append(f"""
                <div class="image-item" data-index="{i}">
                    <img src="{img_src}" alt="Uploaded image {i + 1}">
                    <div class="delete-btn" onclick="deleteImageByIndex({i})" data-index="{i}">×</div>
                </div>
                """)
        except:
            continue

    if not html_items:
        return "<div id='image-display-area'><div class='no-images'>No images yet, please upload images</div></div>"

    html_content = f"""
    <div id='image-display-area'>
        {''.join(html_items)}
    </div>
    """

    return html_content


def handle_file_upload(current_files, new_files):
    """简化的文件上传处理"""
    if not new_files:
        return current_files or [], "Please select images"

    # 确保是列表格式
    current_list = current_files or []
    new_list = new_files if isinstance(new_files, list) else [new_files]

    # 合并并限制数量
    all_files = current_list + [f for f in new_list if f]
    if len(all_files) > 5:
        all_files = all_files[:5]
        message = f"Uploaded {len(all_files)} images (limit reached)"
    else:
        message = f"Uploaded {len(all_files)} images"

    return all_files, message


def process_nano_banana(prompt, uploaded_files, api_key, progress=gr.Progress()):
    """
    处理Nano Banana图像生成的主函数
    """
    # 验证输入
    if not api_key or not api_key.strip():
        return [], "❌ Please enter API Key"

    if not prompt or not prompt.strip():
        return [], "❌ Please enter a prompt"

    # 验证图片列表
    if not uploaded_files or len(uploaded_files) == 0:
        return [], "❌ Please upload at least one image"

    # 确保是列表格式
    file_paths = uploaded_files if isinstance(uploaded_files, list) else [uploaded_files]
    # 验证图片数量
    if len(file_paths) > 5:
        return [], "❌ Maximum 5 images allowed"

    try:
        progress(0.1, desc="📤 Processing images...")

        # 上传文件到 KIE AI 并获取公开URL
        image_urls = []

        for i, file_path in enumerate(file_paths):
            progress(0.1 + (0.3 * i / len(file_paths)), desc=f"📤 Uploading image {i + 1}/{len(file_paths)}...")

            success, result = upload_file_to_kie(file_path, api_key.strip())
            if success:
                image_urls.append(result)
                # 上传成功后删除本地临时文件
                try:
                    if os.path.exists(file_path):
                        # 记录文件所在的目录
                        temp_dir = os.path.dirname(file_path)
                        used_temp_dirs.add(temp_dir)
                        
                        os.remove(file_path)
                except:
                    pass
            else:
                return [], f"❌ Failed to upload image {i + 1}: {result}"

        progress(0.5, desc="🚀 Creating processing task...")

        # 创建任务
        status_code, result = create_task(prompt.strip(), image_urls, api_key.strip())

        if status_code != 200:
            return [], f"❌ Failed to create task: {result}"

        task_id = result
        progress(0.6, desc=f"📋 Task ID: {task_id}")

        progress(0.7, desc="⏳ Processing images, please wait...")

        # 轮询获取结果
        result = get_task_result(task_id, api_key.strip())

        if isinstance(result, tuple) and result[0] == 500:
            return [], f"❌ Task processing failed: {result[1]}"

        progress(0.9, desc="📥 Downloading generated images...")

        # 处理返回的图片URL数组
        generated_images = []
        if isinstance(result, list) and len(result) > 0:
            total_images = len(result)
            for i, img_url in enumerate(result):
                try:
                    progress(0.9 + (0.1 * i / total_images), desc=f"📥 Downloading image {i + 1}/{total_images}...")
                    response = requests.get(img_url, timeout=60)
                    if response.status_code == 200:
                        img = Image.open(BytesIO(response.content))
                        generated_images.append(img)
                except Exception as e:
                    pass

            if generated_images:
                progress(1.0, desc="✅ Complete!")
                return generated_images, f"✅ Successfully generated {len(generated_images)} images!"
            else:
                return [], "❌ Unable to download any generated images"
        else:
            return [], "❌ No generated images received"

    except:
        return [], f"❌ Error occurred during processing"
    finally:
        # 最终清理：删除任何剩余的临时文件
        try:
            for file_path in file_paths:
                if os.path.exists(file_path):
                    # 记录文件所在的目录
                    temp_dir = os.path.dirname(file_path)
                    used_temp_dirs.add(temp_dir)
                    
                    os.remove(file_path)
        except:
            pass


# CSS样式，参考app.py的风格
css = """
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #ffd93d 0%, #ffb347 100%);
    padding: 2.5rem;
    border-radius: 24px;
    margin-bottom: 2.5rem;
    box-shadow: 0 20px 60px rgba(102, 126, 234, 0.25);
}
.logo-text {
    font-size: 3.5rem;
    font-weight: 900;
    color: #2d3436;
    text-align: center;
    margin: 0;
    letter-spacing: -2px;
}
.subtitle {
    color: #2d3436;
    text-align: center;
    font-size: 1rem;
    margin-top: 0.5rem;
    opacity: 0.8;
}
.main-content {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    padding: 2.5rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
}
.mode-indicator {
    background: rgba(255, 255, 255, 0.3);
    backdrop-filter: blur(10px);
    border-radius: 12px;
    padding: 0.5rem 1rem;
    margin-top: 1rem;
    text-align: center;
    font-weight: 600;
    color: #2d3436;
}
.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 1.2rem 2rem !important;
    border-radius: 14px !important;
    text-transform: uppercase;
    letter-spacing: 1px;
    width: 100%;
    margin-top: 1rem !important;
}
.gr-button-primary:hover {
    background: linear-gradient(135deg, #5a67d8 0%, #6b46c1 100%) !important;
}
.gr-input, .gr-textarea {
    background: #ffffff !important;
    border: 1px solid #d1d5db !important;
    border-radius: 8px !important;
    color: #374151 !important;
    font-size: 1rem !important;
    padding: 0.75rem 1rem !important;
}
.gr-input:focus, .gr-textarea:focus {
    border-color: #667eea !important;
    outline: none !important;
}
.gr-form {
    background: transparent !important;
    border: none !important;
}
.gr-panel {
    background: #ffffff !important;
    border: 2px solid #e1e8ed !important;
    border-radius: 16px !important;
    padding: 1.5rem !important;
}
.gr-box {
    border-radius: 14px !important;
    border-color: #e1e8ed !important;
}
label {
    color: #636e72 !important;
    font-weight: 600 !important;
    font-size: 0.85rem !important;
    text-transform: uppercase;
    letter-spacing: 0.5px;
    margin-bottom: 0.5rem !important;
}
.status-text {
    font-family: 'SF Mono', 'Monaco', monospace;
    font-size: 0.95rem;
}
.image-container {
    border-radius: 14px !important;
    overflow: hidden;
    border: 2px solid #e1e8ed !important;
    background: #fafbfc !important;
}
footer {
    display: none !important;
}
.info-box {
    background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%);
    border-radius: 12px;
    padding: 1rem;
    margin-bottom: 1rem;
    border-left: 4px solid #2196f3;
}
.warning-box {
    background: #fff3cd;
    border-radius: 12px;
    padding: 1rem;
    margin-top: 1rem;
    border-left: 4px solid #ffc107;
    color: #856404;
}
/* 简化的图片展示区域 */
.image-upload-container {
    background: white !important;
    border-radius: 16px !important;
    padding: 1.5rem !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.06) !important;
}

/* 自定义图片展示区 - 无边框设计 */
#image-display-area {
    display: flex !important;
    flex-wrap: nowrap !important;
    gap: 12px !important;
    padding: 15px 5px !important;
    overflow-x: auto !important;
    overflow-y: hidden !important;
    min-height: 120px !important;
    max-height: 120px !important;
    align-items: center !important;
    background: transparent !important;
}

/* 无图片时的提示 */
.no-images {
    color: #9ca3af !important;
    font-size: 14px !important;
    text-align: center !important;
    width: 100% !important;
    padding: 20px !important;
}

/* 图片项容器 */
.image-item {
    position: relative !important;
    flex-shrink: 0 !important;
    width: 100px !important;
    height: 100px !important;
    border-radius: 8px !important;
    overflow: hidden !important;
    cursor: pointer !important;
    transition: transform 0.2s ease, box-shadow 0.2s ease !important;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1) !important;
}

.image-item:hover {
    transform: scale(1.05) !important;
    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15) !important;
}

/* 图片样式 */
.image-item img {
    width: 100% !important;
    height: 100% !important;
    object-fit: cover !important;
    border-radius: 8px !important;
    display: block !important;
}

/* 删除按钮 */
.image-item .delete-btn {
    position: absolute !important;
    top: -8px !important;
    right: -8px !important;
    width: 24px !important;
    height: 24px !important;
    background: #ef4444 !important;
    color: white !important;
    border: 2px solid white !important;
    border-radius: 50% !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    font-size: 12px !important;
    font-weight: bold !important;
    cursor: pointer !important;
    opacity: 0 !important;
    transition: all 0.2s ease !important;
    z-index: 10 !important;
    line-height: 1 !important;
}

.image-item:hover .delete-btn {
    opacity: 1 !important;
}

.image-item .delete-btn:hover {
    background: #dc2626 !important;
    transform: scale(1.1) !important;
}

/* 隐藏独立的删除按钮组 */
#delete-buttons-row {
    display: none !important;
}

/* 滚动条样式 */
#image-display-area::-webkit-scrollbar {
    height: 6px !important;
}

#image-display-area::-webkit-scrollbar-track {
    background: rgba(0, 0, 0, 0.05) !important;
    border-radius: 3px !important;
}

#image-display-area::-webkit-scrollbar-thumb {
    background: rgba(0, 0, 0, 0.2) !important;
    border-radius: 3px !important;
}

#image-display-area::-webkit-scrollbar-thumb:hover {
    background: rgba(0, 0, 0, 0.3) !important;
}

/* 更小的上传文件框 */
.gr-file {
    background: white !important;
    border: 1px dashed #d1d5db !important;
    border-radius: 4px !important;
    padding: 0.1rem 0.3rem !important;
    font-size: 0.65rem !important;
    min-height: 16px !important;
    height: 16px !important;
    max-width: 60px !important;
}

.gr-file:hover {
    border-color: #667eea !important;
    background: #f9fafb !important;
}

.gr-file .wrap {
    min-height: 14px !important;
    padding: 0 !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
}

.gr-file .wrap > div {
    font-size: 0.65rem !important;
    color: #6b7280 !important;
    line-height: 1 !important;
}


/* 输出标题样式 */
.output-title {
    margin-bottom: 0.5rem !important;
    margin-top: 0 !important;
}

.output-title h3 {
    margin: 0 !important;
    padding: 0 !important;
    font-size: 1.1rem !important;
    color: #374151 !important;
    font-weight: 600 !important;
}

/* 输出图片画廊 */
#output-gallery {
    border: 2px solid #e1e8ed !important;
    border-radius: 14px !important;
    padding: 1rem !important;
    background: #fafbfc !important;
    height: 500px !important;  /* 固定高度 */
    overflow: hidden !important;  /* 隐藏外层滚动条 */
    position: relative !important;
}

/* 输出画廊内部容器控制 */
#output-gallery > div {
    height: 100% !important;
    overflow: hidden !important;
}

#output-gallery .gr-gallery {
    height: 100% !important;
    position: relative !important;
}

/* 输出画廊滚动容器 - 只允许垂直滚动 */
#output-gallery .gr-gallery-container {
    position: absolute !important;
    top: 0 !important;
    left: 0 !important;
    right: 0 !important;
    bottom: 0 !important;
    overflow-y: auto !important;
    overflow-x: hidden !important;
    gap: 1rem !important;
    padding: 5px !important;
}

/* 美化输出画廊的滚动条 */
#output-gallery .gr-gallery-container::-webkit-scrollbar {
    width: 8px !important;
}

#output-gallery .gr-gallery-container::-webkit-scrollbar-track {
    background: rgba(0, 0, 0, 0.05) !important;
    border-radius: 4px !important;
}

#output-gallery .gr-gallery-container::-webkit-scrollbar-thumb {
    background: rgba(0, 0, 0, 0.2) !important;
    border-radius: 4px !important;
}

#output-gallery .gr-gallery-container::-webkit-scrollbar-thumb:hover {
    background: rgba(0, 0, 0, 0.3) !important;
}

.gr-group {
    background: #f8f9fa !important;
    border-radius: 14px !important;
    padding: 1rem !important;
    margin-bottom: 1rem !important;
}
"""

# 创建Gradio界面
with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    with gr.Column(elem_classes="header-container"):
        gr.HTML("""
            <h1 class="logo-text">🍌 Nano Banana API Free Online Test</h1>
            <p class="subtitle">Powered by Google’s Official Gemini 2.5 Flash Image Model</p>
            <div class="mode-indicator">
                💡 Nano Banana AI makes editing smarter — upload up to 5 photos, refine details, and maintain subject consistency. Free to try: 80 credits = 20 images.
            </div>
        """)

    with gr.Column(elem_classes="main-content"):
        # 信息提示框
        gr.HTML("""
            <div class="info-box">
                <strong>Usage Instructions:</strong><br>
                • Get your Nano Banana API Key   <a href="https://kie.ai/nano-banana" target="_blank">👉 here 👈</a><br>
                • Add your image and enter a prompt to generate or edit.<br>
                • Upload 1–5 reference images (Max 10MB each, formats: JPG, PNG, WebP)<br>
                • Click Generate and wait ~1–2 minutes for processing
            </div>
        """)

        with gr.Row(equal_height=True):
            # 左侧 - 输入区域
            with gr.Column(scale=1):
                # API Key输入
                api_key = gr.Textbox(
                    label="API Key",
                    placeholder="Please enter your KIE AI API Key",
                    type="password",
                    elem_classes="api-key-input"
                )

                # 提示词输入
                prompt = gr.Textbox(
                    label="Editing Prompt",
                    placeholder="Describe the image effect you want, e.g.: Convert image to cartoon style...",
                    lines=3,
                    value="Transform the image into a dreamy oil painting style with vibrant colors and clear brushstrokes",
                    elem_classes="prompt-input"
                )

                # 重新设计的图片上传区域
                with gr.Group(elem_classes="image-upload-container"):
                    gr.Markdown("### 📸 Image Upload")

                    # 自定义图片展示区（无边框）
                    image_display = gr.HTML(
                        value="<div id='image-display-area'><div class='no-images'>No images yet, please upload images</div></div>",
                        elem_id="image-display"
                    )

                    # 简单的上传按钮
                    file_upload = gr.File(
                        show_label=False,
                        file_count="multiple",
                        file_types=["image"],
                        type="filepath",
                        height=120,  # 缩小高度
                    )

                # 添加删除按钮组（在图片上传区域内）- 通过CSS隐藏而不是visible=False
                with gr.Row(elem_id="delete-buttons-row"):
                    delete_label = gr.Markdown("**Delete Images:**", visible=True, elem_id="delete-label")
                    delete_buttons = []
                    for i in range(5):  # 最多支持5张图片
                        btn = gr.Button(f"Delete {i + 1}", visible=True, size="sm", elem_id=f"delete-btn-{i}")
                        delete_buttons.append(btn)

                # 生成按钮
                generate_btn = gr.Button(
                    "🚀 Start Generation",
                    variant="primary",
                    size="lg"
                )

            # 右侧 - 输出区域
            with gr.Column(scale=1):
                # 添加输出标题
                gr.Markdown("### 🎨 Generation Results", elem_classes="output-title")

                # 输出图片画廊
                output_gallery = gr.Gallery(
                    show_label=False,
                    elem_id="output-gallery",
                    columns=2,
                    rows=None,
                    object_fit="contain",
                    height=500,
                    preview=True,
                    container=True
                )

                # 状态信息
                status = gr.Textbox(
                    label="Processing Status",
                    interactive=False,
                    lines=2,
                    value="Ready, please upload images and enter a prompt..."
                )

        # 示例
        gr.Examples(
            examples=[
                ["Convert image to cartoon anime style", None],
                ["Apply Van Gogh's starry night style", None],
                ["Convert to black and white sketch", None],
                ["Add neon light effects, cyberpunk style", None],
                ["Convert to watercolor painting style with soft tones", None],
            ],
            inputs=[prompt, api_key],
            label="Prompt Examples"
        )

    # 状态变量
    current_files = gr.State([])


    # 更新的事件处理函数
    def on_file_upload(current_paths, new_files):
        """处理文件上传并更新HTML显示"""
        updated_files, message = handle_file_upload(current_paths, new_files)
        html_content = create_image_html(updated_files)
        return updated_files, gr.update(value=None), html_content


    # 为每个删除按钮绑定事件
    def create_delete_handler(index):
        def delete_image_at_index(current_paths):
            if not current_paths or index >= len(current_paths):
                return current_paths, create_image_html(current_paths), *update_delete_buttons(current_paths)

            # 删除指定索引的图片
            updated_files = current_paths[:index] + current_paths[index + 1:]
            return updated_files, create_image_html(updated_files), *update_delete_buttons(updated_files)

        return delete_image_at_index


    def update_delete_buttons(file_paths):
        """更新删除按钮 - 保持所有按钮visible=True，通过CSS控制显示"""
        updates = []
        updates.append(gr.update())  # 标签保持不变

        for i in range(5):
            if i < len(file_paths):
                updates.append(gr.update(value=f"Delete Image {i + 1}"))
            else:
                updates.append(gr.update())  # 保持不变
        return updates


    # 绑定删除按钮事件
    outputs_list = [current_files, image_display, delete_label] + delete_buttons

    for i, btn in enumerate(delete_buttons):
        btn.click(
            fn=create_delete_handler(i),
            inputs=[current_files],
            outputs=outputs_list
        )


    # 更新文件上传事件，同时更新删除按钮
    def on_file_upload_with_buttons(current_paths, new_files):
        updated_files, message = handle_file_upload(current_paths, new_files)
        html_content = create_image_html(updated_files)
        button_updates = update_delete_buttons(updated_files)
        return updated_files, gr.update(value=None), html_content, *button_updates


    file_upload.upload(
        fn=on_file_upload_with_buttons,
        inputs=[current_files, file_upload],
        outputs=[current_files, file_upload, image_display, delete_label] + delete_buttons
    )


    # 生成按钮事件
    def prepare_and_generate(prompt, paths, api_key, progress=gr.Progress()):
        """生成图片的主函数"""
        if not paths:
            return [], "❌ Please upload at least one image"
        return process_nano_banana(prompt, paths, api_key, progress)


    generate_btn.click(
        fn=prepare_and_generate,
        inputs=[prompt, current_files, api_key],
        outputs=[output_gallery, status]
    )

    # 添加JavaScript代码来连接图片上的删除按钮和隐藏的独立删除按钮
    demo.load(None, None, None, js="""
    () => {
        // 创建全局删除函数
        window.deleteImageByIndex = function(index) {
            // Gradio的elem_id设置在包装器上，需要找到内部的button
            let deleteBtn = null;

            // 方法1: 通过ID找到包装器，然后找内部的button
            const wrapper = document.getElementById(`delete-btn-${index}`);
            if (wrapper) {
                deleteBtn = wrapper.querySelector('button');
            }

            // 方法2: 如果方法1失败，查找所有按钮并通过文本内容匹配
            if (!deleteBtn) {
                const allButtons = document.querySelectorAll('button');
                for (let btn of allButtons) {
                    if (btn.textContent.includes(`Delete Image ${index + 1}`)) {
                        deleteBtn = btn;
                        break;
                    }
                }
            }

            if (deleteBtn) {
                deleteBtn.click();
            } else {
                // 调试信息
                document.querySelectorAll('[id^="delete-btn-"]').forEach(elem => {
                    console.log(elem.id, elem.tagName, elem.querySelector('button'));
                });
            }
        };
    }
    """)

# 启动应用
if __name__ == "__main__":
    app_instance = demo

    # 启动定时清理任务
    start_cleanup_scheduler()
    
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860  # 使用不同的端口避免冲突
    )