app.py

import gradio as gr
import requests
import json
import os
import subprocess
import uuid
import time
import cv2
from typing import Optional, List
import numpy as np
from datetime import datetime, timedelta
from collections import defaultdict

os.environ["SPACES_QUEUE_ENABLED"] = "true"

# 后端API配置（可配置化）
BACKEND_URL = os.getenv("BACKEND_URL", "http://47.95.6.204:51001/")
API_ENDPOINTS = {
    "submit_task": f"{BACKEND_URL}/predict/video",
    "query_status": f"{BACKEND_URL}/predict/task",
    "terminate_task": f"{BACKEND_URL}/predict/terminate"
}

# 模拟场景配置
SCENE_CONFIGS = {
    "scene_1": {
        "description": "scene_1",
        "objects": ["milk carton", "ceramic bowl", "mug"],
        "preview_image": "assets/scene_1.png"
    },
}

# 可用模型列表
MODEL_CHOICES = [
    "gr1",
    # "GR00T-N1",
    # "GR00T-1.5",
    # "Pi0",
    # "DP+CLIP",
    # "AcT+CLIP"
]

###############################################################################

SESSION_TASKS = {}
IP_REQUEST_RECORDS = defaultdict(list)
IP_LIMIT = 5  # 每分钟最多请求次数

def is_request_allowed(ip: str) -> bool:
    now = datetime.now()
    IP_REQUEST_RECORDS[ip] = [t for t in IP_REQUEST_RECORDS[ip] if now - t < timedelta(minutes=1)]
    if len(IP_REQUEST_RECORDS[ip]) < IP_LIMIT:
        IP_REQUEST_RECORDS[ip].append(now)
        return True
    return False
###############################################################################


# 日志文件路径
LOG_DIR = "/opt/gradio-frontend/logs"
os.makedirs(LOG_DIR, exist_ok=True)
ACCESS_LOG = os.path.join(LOG_DIR, "access.log")
SUBMISSION_LOG = os.path.join(LOG_DIR, "submissions.log")

def log_access(user_ip: str = None, user_agent: str = None):
    """记录用户访问日志"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = {
        "timestamp": timestamp,
        "type": "access",
        "user_ip": user_ip or "unknown",
        "user_agent": user_agent or "unknown"
    }
    
    with open(ACCESS_LOG, "a") as f:
        f.write(json.dumps(log_entry) + "\n")

def log_submission(scene: str, prompt: str, model: str, max_step: int, user: str = "anonymous", res: str = "unknown"):
    """记录用户提交日志"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    log_entry = {
        "timestamp": timestamp,
        "type": "submission",
        "user": user,
        "scene": scene,
        "prompt": prompt,
        "model": model,
        "max_step": str(max_step),
        "res": res
    }
    
    with open(SUBMISSION_LOG, "a") as f:
        f.write(json.dumps(log_entry) + "\n")

def read_logs(log_type: str = "all", max_entries: int = 50) -> list:
    """读取日志文件"""
    logs = []
    
    if log_type in ["all", "access"]:
        try:
            with open(ACCESS_LOG, "r") as f:
                for line in f:
                    logs.append(json.loads(line.strip()))
        except FileNotFoundError:
            pass
    
    if log_type in ["all", "submission"]:
        try:
            with open(SUBMISSION_LOG, "r") as f:
                for line in f:
                    logs.append(json.loads(line.strip()))
        except FileNotFoundError:
            pass
    
    # 按时间戳排序，最新的在前
    logs.sort(key=lambda x: x["timestamp"], reverse=True)
    return logs[:max_entries]

def format_logs_for_display(logs: list) -> str:
    """格式化日志用于显示"""
    if not logs:
        return "暂无日志记录"
    
    markdown = "### 系统访问日志\n\n"
    markdown += "| 时间 | 类型 | 用户/IP | 详细信息 |\n"
    markdown += "|------|------|---------|----------|\n"
    
    for log in logs:
        timestamp = log.get("timestamp", "unknown")
        log_type = "访问" if log.get("type") == "access" else "提交"
        
        if log_type == "访问":
            user = log.get("user_ip", "unknown")
            details = f"User-Agent: {log.get('user_agent', 'unknown')}"
        else:
            user = log.get("user", "anonymous")
            result = log.get('res', 'unknown')
            if result != "success": 
                if len(result) > 40:  # Adjust this threshold as needed
                    result = f"{result[:20]}...{result[-20:]}"
            details = f"场景: {log.get('scene', 'unknown')}, 指令: {log.get('prompt', '')}, 模型: {log.get('model', 'unknown')}, max step: {log.get('max_step', '300')}, result: {result}"
        
        markdown += f"| {timestamp} | {log_type} | {user} | {details} |\n"
    
    return markdown



###############################################################################


def stream_simulation_results(result_folder: str, task_id: str, fps: int = 30):
    """
    流式输出仿真结果，同时监控图片文件夹和后端任务状态
    
    参数:
        result_folder: 包含生成图片的文件夹路径
        task_id: 后端任务ID用于状态查询
        fps: 输出视频的帧率
        
    生成:
        生成的视频文件路径 (分段输出)
    """
    # 初始化变量
    result_folder = os.path.join(result_folder, "image")
    os.makedirs(result_folder, exist_ok=True)
    frame_buffer: List[np.ndarray] = []
    frames_per_segment = fps * 2  # 每2秒60帧
    processed_files = set()
    width, height = 0, 0
    last_status_check = 0
    status_check_interval = 5  # 每5秒检查一次后端状态
    max_time = 240

    while max_time > 0:
        max_time -= 1
        current_time = time.time()
        
        # 定期检查后端状态
        if current_time - last_status_check > status_check_interval:
            status = get_task_status(task_id)
            print("status: ", status)
            if status.get("status") == "completed":
                # 确保处理完所有已生成的图片
                process_remaining_images(result_folder, processed_files, frame_buffer)
                if frame_buffer:
                    yield create_video_segment(frame_buffer, fps, width, height)
                break
            elif status.get("status") == "failed":
                raise gr.Error(f"任务执行失败: {status.get('result', '未知错误')}")
            elif status.get("status") == "terminated":
                break
            last_status_check = current_time

        # 处理新生成的图片
        current_files = sorted(
            [f for f in os.listdir(result_folder) 
             if f.lower().endswith(('.png', '.jpg', '.jpeg'))],
            key=lambda x: os.path.splitext(x)[0]  # 按文件名排序
        )
        
        new_files = [f for f in current_files if f not in processed_files]
        has_new_frames = False
        
        for filename in new_files:
            try:
                img_path = os.path.join(result_folder, filename)
                frame = cv2.imread(img_path)
                if frame is not None:
                    if width == 0:  # 第一次获取图像尺寸
                        height, width = frame.shape[:2]
                    
                    frame_buffer.append(frame)
                    processed_files.add(filename)
                    has_new_frames = True
            except Exception as e:
                print(f"Error processing {filename}: {e}")

        # 如果有新帧且积累够60帧，输出视频片段
        if has_new_frames and len(frame_buffer) >= frames_per_segment:
            segment_frames = frame_buffer[:frames_per_segment]
            frame_buffer = frame_buffer[frames_per_segment:]
            yield create_video_segment(segment_frames, fps, width, height)

        time.sleep(1)  # 避免过于频繁检查
    
    if max_time <= 0:
        raise gr.Error("timeout 240s")

def create_video_segment(frames: List[np.ndarray], fps: int, width: int, height: int) -> str:
    """创建视频片段"""
    os.makedirs("/opt/gradio_demo/tasks/video_chunk", exist_ok=True)
    segment_name = f"/opt/gradio_demo/tasks/video_chunk/output_{uuid.uuid4()}.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(segment_name, fourcc, fps, (width, height))
    
    for frame in frames:
        out.write(frame)
    out.release()
    
    return segment_name

def process_remaining_images(result_folder: str, processed_files: set, frame_buffer: List[np.ndarray]):
    """处理剩余的图片"""
    current_files = sorted(
        [f for f in os.listdir(result_folder) 
         if f.lower().endswith(('.png', '.jpg', '.jpeg'))],
        key=lambda x: os.path.splitext(x)[0]
    )
    
    new_files = [f for f in current_files if f not in processed_files]
    
    for filename in new_files:
        try:
            img_path = os.path.join(result_folder, filename)
            frame = cv2.imread(img_path)
            if frame is not None:
                frame_buffer.append(frame)
                processed_files.add(filename)
        except Exception as e:
            print(f"Error processing remaining {filename}: {e}")




###############################################################################



def submit_to_backend(
    scene: str,
    prompt: str,
    model: str,
    max_step: int,
    user: str = "Gradio-user",
) -> dict:
    job_id = str(uuid.uuid4())

    data = {
        "scene_type": scene,
        "instruction": prompt,
        "model_type": model,
        "max_step": str(max_step)
    }
    
    payload = {
        "user": user,
        "task": "robot_manipulation",
        "job_id": job_id,
        "data": json.dumps(data) 
    }
    
    try:
        headers = {"Content-Type": "application/json"}
        response = requests.post(
            API_ENDPOINTS["submit_task"],
            json=payload,
            headers=headers,
            timeout=10
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "message": str(e)}

def get_task_status(task_id: str) -> dict:
    """
    查询任务状态
    """
    try:
        response = requests.get(
            f"{API_ENDPOINTS['query_status']}/{task_id}",
            timeout=5
        )
        return response.json()
    except Exception as e:
        return {"status": "error", "message": str(e)}

def terminate_task(task_id: str) -> Optional[dict]:
    """
    终止任务
    """
    try:
        response = requests.post(
            f"{API_ENDPOINTS['terminate_task']}/{task_id}",
            timeout=3
        )
        return response.json()
    except Exception as e:
        print(f"Error terminate task: {e}")
        return None

def convert_to_h264(video_path):
    """
    将视频转换为 H.264 编码的 MP4 格式
    生成新文件路径在原路径基础上添加 _h264 后缀）
    """
    base, ext = os.path.splitext(video_path)
    video_path_h264 = f"{base}_h264.mp4"
    
    try:
        # 构建 FFmpeg 命令
        ffmpeg_cmd = [
            "ffmpeg",
            "-i", video_path,
            "-c:v", "libx264",
            "-preset", "slow",
            "-crf", "23",
            "-c:a", "aac",
            "-movflags", "+faststart",
            video_path_h264
        ]
        
        # 执行 FFmpeg 命令
        subprocess.run(ffmpeg_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        # 检查输出文件是否存在
        if not os.path.exists(video_path_h264):
            raise FileNotFoundError(f"H.264 编码文件未生成: {video_path_h264}")
            
        return video_path_h264
        
    except subprocess.CalledProcessError as e:
        raise gr.Error(f"FFmpeg 转换失败: {e.stderr}")
    except Exception as e:
        raise gr.Error(f"转换过程中发生错误: {str(e)}")

def run_simulation(
    scene: str,
    prompt: str,
    model: str,
    max_step: int,
    history: list,
    request: gr.Request
) -> dict:
    """运行仿真并更新历史记录"""
    # 获取当前时间
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    scene_desc = SCENE_CONFIGS.get(scene, {}).get("description", scene)
    
    # 记录用户提交
    user_ip = request.client.host if request else "unknown"
    session_id = request.session_hash

    if not is_request_allowed(user_ip):
        log_submission(scene, prompt, model, max_step, user_ip, "IP blocked temporarily")
        raise gr.Error("Too many requests from this IP. Please wait and try again one minute later.")
    
    # 提交任务到后端
    submission_result = submit_to_backend(scene, prompt, model, max_step, user_ip)
    print("submission_result: ", submission_result)

    if submission_result.get("status") != "pending":
        log_submission(scene, prompt, model, max_step, user_ip, "Submission failed")
        raise gr.Error(f"Submission failed: {submission_result.get('message', 'unknown issue')}")
    
    try:
        task_id = submission_result["task_id"]
        SESSION_TASKS[session_id] = task_id

        gr.Info(f"Simulation started, task_id: {task_id}")
        time.sleep(5)
        # 获取任务状态
        status = get_task_status(task_id)
        print("first status: ", status)
        result_folder = status.get("result", "")
    except Exception as e:
        log_submission(scene, prompt, model, max_step, user_ip, str(e))
        raise gr.Error(f"error occurred when parsing submission result from backend: {str(e)}")


    if not os.path.exists(result_folder):
        log_submission(scene, prompt, model, max_step, user_ip, "Result folder provided by backend doesn't exist")
        raise gr.Error(f"Result folder provided by backend doesn't exist: <PATH>{result_folder}")
    
    
    # 流式输出视频片段
    try:
        for video_path in stream_simulation_results(result_folder, task_id):
            if video_path:
                yield video_path, history
    except Exception as e:
        log_submission(scene, prompt, model, max_step, user_ip, str(e))
        raise gr.Error(f"Error while streaming: {str(e)}")

    # 获取任务状态
    status = get_task_status(task_id)
    print("status: ", status)
    if status.get("status") == "completed":
        video_path = os.path.join(status.get("result"), "manipulation.mp4")
        print("video_path: ", video_path)
        video_path = convert_to_h264(video_path)

        # 创建新的历史记录条目
        new_entry = {
            "timestamp": timestamp,
            "scene": scene,
            "model": model,
            "prompt": prompt,
            "max_step": max_step,
            "video_path": video_path,
            "task_id": task_id
        }
        
        # 将新条目添加到历史记录顶部
        updated_history = history + [new_entry]
        
        # 限制历史记录数量，避免内存问题
        if len(updated_history) > 10:
            updated_history = updated_history[:10]
        
        print("updated_history:", updated_history)
        log_submission(scene, prompt, model, max_step, user_ip, "success")
        gr.Info("Simulation completed successfully!")
        yield None, updated_history

    elif status.get("status") == "failed":
        log_submission(scene, prompt, model, max_step, user_ip, status.get('result', 'backend error'))
        raise gr.Error(f"Task execution failed: {status.get('result', 'backend unknown issue')}")
        yield None, history
    
    elif status.get("status") == "terminated":
        log_submission(scene, prompt, model, max_step, user_ip, "user end terminated")
        yield None, history

    else:
        log_submission(scene, prompt, model, max_step, user_ip, "missing task's status from backend (Pending?)")
        raise gr.Error("missing task's status from backend (Pending?)")
        yield None, history


###############################################################################

    
def update_history_display(history: list) -> list:
    """更新历史记录显示"""
    print("更新历史记录显示")
    updates = []
    
    for i in range(10):
        if i < len(history):  # 如果有历史记录，更新对应槽位
            entry = history[i]
            updates.extend([
                gr.update(visible=True),  # 更新 Column 可见性
                gr.update(visible=True, label=f"# {i+1} | {entry['scene']} | {entry['model']} | {entry['prompt']}", open=(i+1==len(history))),  # 更新 Accordion
                gr.update(value=entry['video_path'], visible=True, autoplay=False),  # 更新 Video
                gr.update(value=f"{entry['timestamp']}")  # 更新详细 Markdown
            ])
        else:  # 如果没有历史记录，隐藏槽位
            updates.extend([
                gr.update(visible=False),  # 隐藏 Column
                gr.update(visible=False),  # 隐藏 Accordion
                gr.update(value=None, visible=False),  # 清空 Video
                gr.update(value="")  # 清空详细 Markdown
            ])
    print("更新完成！")
    return updates

def update_scene_display(scene: str) -> tuple[str, Optional[str]]:
    """更新场景描述和预览图"""
    config = SCENE_CONFIGS.get(scene, {})
    desc = config.get("description", "No description")
    objects = ", ".join(config.get("objects", []))
    image = config.get("preview_image", None)
    
    markdown = f"**{desc}**  \nObjects in this scene: {objects}"
    return markdown, image

def update_log_display():
    """更新日志显示"""
    logs = read_logs()
    return format_logs_for_display(logs)

###############################################################################


def cleanup_session(request: gr.Request):
    session_id = request.session_hash
    task_id = SESSION_TASKS.pop(session_id, None)
    
    if task_id:
        try:
            status = get_task_status(task_id)
            print("clean up check status: ", status)
            if status.get("status") == "pending":
                res = terminate_task(task_id)
                if res.get("status") == "success":
                    print(f"已终止任务 {task_id}")
                else:
                    print(f"终止任务失败 {task_id}: {res.get('status', 'unknown issue')}")
        except Exception as e:
            print(f"终止任务失败 {task_id}: {e}")



###############################################################################

header_html = """
<div style="display: flex; justify-content: space-between; align-items: center; width: 100%; margin-bottom: 20px; padding: 20px; background: linear-gradient(135deg, #528bdb 0%, #a7b5d0 100%); border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);">
    <div style="display: flex; align-items: center;">
        <img src="https://www.shlab.org.cn/static/img/index_14.685f6559.png" alt="Institution Logo" style="height: 60px; margin-right: 20px;">
        <div>
            <h1 style="margin: 0; color: #2c3e50; font-weight: 600;">🤖 InternManip Model Inference Demo</h1>
            <p style="margin: 4px 0 0 0; color: #5d6d7e; font-size: 0.9em;">Model trained on InternManip framework</p>
        </div>
    </div>
    <div style="display: flex; gap: 15px; align-items: center;">
        <a href="https://github.com/InternRobotics" target="_blank" style="text-decoration: none; transition: transform 0.2s;" onmouseover="this.style.transform='scale(1.1)'" onmouseout="this.style.transform='scale(1)'">
            <img src="https://github.githubassets.com/images/modules/logos_page/GitHub-Mark.png" alt="GitHub" style="height: 30px;">
        </a>
        <a href="https://huggingface.co/InternRobotics" target="_blank" style="text-decoration: none; transition: transform 0.2s;" onmouseover="this.style.transform='scale(1.1)'" onmouseout="this.style.transform='scale(1)'">
            <img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" alt="HuggingFace" style="height: 30px;">
        </a>
        <a href="https://huggingface.co/spaces/OpenRobotLab/InternNav-eval-demo" target="_blank">
            <button style="padding: 8px 15px; background: #3498db; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: 500; transition: all 0.2s;" 
                    onmouseover="this.style.backgroundColor='#2980b9'; this.style.transform='scale(1.05)'" 
                    onmouseout="this.style.backgroundColor='#3498db'; this.style.transform='scale(1)'">
                Go to InternNav Demo
            </button>
        </a>
    </div>
</div>
"""

###############################################################################


# 自定义CSS样式
custom_css = """
#simulation-panel {
    border-radius: 8px;
    padding: 20px;
    background: #f9f9f9;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
#result-panel {
    border-radius: 8px;
    padding: 20px;
    background: #f0f8ff;
}
.dark #simulation-panel { background: #2a2a2a; }
.dark #result-panel { background: #1a2a3a; }

.history-container {
    max-height: 600px;
    overflow-y: auto;
    margin-top: 20px;
}

.history-accordion {
    margin-bottom: 10px;
}

.logs-container {
    max-height: 500px;
    overflow-y: auto;
    margin-top: 20px;
    padding: 15px;
    background: #f5f5f5;
    border-radius: 8px;
}

.dark .logs-container {
    background: #2a2a2a;
}

.log-table {
    width: 100%;
    border-collapse: collapse;
}

.log-table th, .log-table td {
    padding: 8px 12px;
    border: 1px solid #ddd;
    text-align: left;
}

.dark .log-table th, .dark .log-table td {
    border-color: #444;
}
"""

# 创建Gradio界面
with gr.Blocks(title="InternManip Model Inference Demo", css=custom_css) as demo:
    gr.HTML(header_html)

    # # 标题和描述
    # gr.Markdown("""
    # # 🤖 InternManip Model Inference Demo
    # ### Model trained on InternManip framework
    # """)
    
    # 存储历史记录的组件变量
    history_state = gr.State([])

    with gr.Row():
        # 左侧控制面板
        with gr.Column(elem_id="simulation-panel"):
            gr.Markdown("### Simulation Settings")
            
            # 场景选择
            scene_dropdown = gr.Dropdown(
                label="Choose a scene",
                choices=list(SCENE_CONFIGS.keys()),
                value="scene_1",
                interactive=True
            )
            
            # 场景描述预览
            scene_description = gr.Markdown("")
            scene_preview = gr.Image(
                label="Scene Preview",
                elem_classes=["scene-preview"],
                interactive=False
            )
            
            scene_dropdown.change(
                update_scene_display,
                inputs=scene_dropdown,
                outputs=[scene_description, scene_preview]
            )
            
            # 操作指令输入
            prompt_input = gr.Textbox(
                label="Manipulation Prompt",
                value="Move the milk carton to the top of the ceramic bowl.",
                placeholder="Example: 'Move the milk carton to the top of the ceramic bowl.'",
                lines=2,
                max_lines=4
            )
            
            # 模型选择
            model_dropdown = gr.Dropdown(
                label="Chose a pretrained model",
                choices=MODEL_CHOICES,
                value=MODEL_CHOICES[0]
            )

            with gr.Accordion("Advance Settings", open=False):
                max_steps = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=300,
                    step=10,
                    label="Max Steps"
                )
            
            # 提交按钮
            submit_btn = gr.Button("Apply and Start Simulation", variant="primary")
        
        # 右侧结果面板
        with gr.Column(elem_id="result-panel"):
            gr.Markdown("### Result")

            # progress_instruction = gr.Markdown("### Please click the botton on the left column to start.")

            # 视频输出
            video_output = gr.Video(
                label="Live",
                interactive=False,
                format="mp4",
                autoplay=True,
                streaming=True
            )
            
            # 历史记录显示区域
            with gr.Column() as history_container:
                gr.Markdown("### History")
                gr.Markdown("#### History will be reset after refresh")
                
                # 预创建10个历史记录槽位
                history_slots = []
                for i in range(10):
                    with gr.Column(visible=False) as slot:
                        with gr.Accordion(visible=False, open=False) as accordion:
                            video = gr.Video(interactive=False)  # 用于播放视频
                            detail_md = gr.Markdown()  # 用于显示详细信息
                    history_slots.append((slot, accordion, video, detail_md))  # 存储所有相关组件
    
    # 添加日志显示区域
    with gr.Accordion("查看系统访问日志(DEV ONLY)", open=False):
        logs_display = gr.Markdown()
        refresh_logs_btn = gr.Button("刷新日志", variant="secondary")
        
        refresh_logs_btn.click(
            update_log_display,
            outputs=logs_display
        )
    
    # 示例
    gr.Examples(
        examples=[
            ["scene_1", "Move the milk carton to the top of the ceramic bowl.", "gr1", 300],
        ],
        inputs=[scene_dropdown, prompt_input, model_dropdown, max_steps],
        label="Examples"
    )
    
    # 提交处理
    submit_btn.click(
        fn=run_simulation,
        inputs=[scene_dropdown, prompt_input, model_dropdown, max_steps, history_state],
        outputs=[video_output, history_state],
        queue=True
    ).then(
        fn=update_history_display,
        inputs=history_state,
        outputs=[comp for slot in history_slots for comp in slot],
        queue=True
    ).then(
        fn=update_log_display,
        outputs=logs_display
    )
    
    # 初始化场景描述和日志
    demo.load(
        fn=lambda: update_scene_display("scene_1"),
        outputs=[scene_description, scene_preview]
    ).then(
        fn=update_log_display,
        outputs=logs_display
    )
    
    # 记录访问
    def record_access(request: gr.Request):
        user_ip = request.client.host if request else "unknown"
        user_agent = request.headers.get("user-agent", "unknown")
        log_access(user_ip, user_agent)
        return update_log_display()
    
    demo.load(
        fn=record_access,
        inputs=None,
        outputs=logs_display,
        queue=False
    )
    
    demo.queue(default_concurrency_limit=8)

    demo.unload(fn=cleanup_session)


# 启动应用
if __name__ == "__main__":
    demo.launch()