# 安装atomgit SDK
!pip install -U atomgit -i https://mirrors.huaweicloud.com/repository/pypi/simple

# restart the kernel
get_ipython().kernel.do_shutdown(restart=True)

# 下载模型
from atomgit_hub import snapshot_download
snapshot_download(
    "hf_mirrors/Qwen/Qwen2.5-Omni-7B",
    local_dir = '/opt/atomgit/Qwen2.5-Omni-7B'
)

%%bash
PORT=18004
MODEL_DIR=/opt/atomgit/Qwen2.5-Omni-7B
LOG=/opt/atomgit/vllm.log
PID_FILE=/opt/atomgit/vllm.pid
LOCK_FILE=/opt/atomgit/vllm_start.lock

PORT_HEX=$(printf '%04X\n' $PORT)

echo "🔍 检查 vLLM 状态..."

# -------------------------
# 防止“正在启动中”
# -------------------------
if [ -f "$LOCK_FILE" ]; then
  echo "⏳ vLLM 正在启动中（LOCK 存在），请稍后..."
  exit 0
fi

# 创建启动锁
touch "$LOCK_FILE"

# -------------------------
# PID 存活检查
# -------------------------
if [ -f "$PID_FILE" ]; then
  PID=$(cat "$PID_FILE")
  if ps -p "$PID" > /dev/null 2>&1; then
    echo "❌ vLLM 已运行，PID=$PID"
    rm -f "$LOCK_FILE"
    exit 0
  else
    echo "⚠️ 旧 PID 已失效，清理"
    rm -f "$PID_FILE"
  fi
fi

# -------------------------
# 端口检查（防已启动成功）
# -------------------------
if grep -i ":$PORT_HEX" /proc/net/tcp > /dev/null; then
  echo "❌ vLLM 已在运行（端口 $PORT 占用）"
  rm -f "$LOCK_FILE"
  exit 0
fi

# -------------------------
# 启动 vLLM
# -------------------------
echo "🚀 启动 vLLM..."

nohup vllm serve "$MODEL_DIR" \
  --served-model-name Qwen2.5-Omni-7B \
  --host 0.0.0.0 \
  --port $PORT \
  --tensor-parallel-size 1 \
  --dtype bfloat16 \
  --compilation-config '{"custom_ops":["none", "+rms_norm", "+rotary_embedding"]}' \
  --max-num-seqs 4 \
  --max-model-len 16000 \
  --gpu-memory-utilization 0.8 \
  > "$LOG" 2>&1 &

echo $! > "$PID_FILE"

echo "⏳ vLLM 启动中，PID=$(cat $PID_FILE)"

# -------------------------
# 延迟释放锁（关键）
# -------------------------
sleep 2
rm -f "$LOCK_FILE"

echo "✅ 启动流程结束"

# 测试数据准备
!mkdir -p static
!curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/image.jpg -o /opt/atomgit/static/image.jpg
!curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/audio.wav -o /opt/atomgit/static/audio.wav
!curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/video.mp4 -o /opt/atomgit/static/video.mp4

# call_api method
import base64
import requests
import json
import os
from requests.exceptions import ConnectionError, Timeout

# API 地址
BASE_URL = "http://localhost:18004/v1/chat/completions"
IMAGE_PATH = "/opt/atomgit/static/image.jpg"
VIDEO_PATH = "/opt/atomgit/static/video.mp4"
AUDIO_PATH = "/opt/atomgit/static/audio.wav"
VLLM_LOG_PATH = "/opt/atomgit/vllm.log"

def encode_file_to_base64(file_path: str) -> str:
    """将文件编码为base64字符串（去掉换行）"""
    with open(file_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8').replace('\n', '')

def call_api(payload, url: str = BASE_URL, timeout: int = 10):
    """
    调用 vLLM OpenAI 兼容接口
    - 捕获端口不存在 / 服务未启动
    - 统一返回 JSON
    """
    headers = {"Content-Type": "application/json"}

    try:
        response = requests.post(
            url,
            headers=headers,
            json=payload,
            timeout=timeout
        )

        print(f"状态码: {response.status_code}")

        # 正常返回
        if response.status_code == 200:
            return response.json()

        # vLLM 返回的业务错误（例如模型未加载）
        try:
            return response.json()
        except Exception:
            return {
                "error": f"vllm 返回非 JSON 响应，状态码 {response.status_code}",
                "raw_response": response.text
            }

    except ConnectionError:
        # 端口不存在 / 服务没起
        return {
            "error": "vllm serve 暂未启动，稍后重试，启动日志详情见：/opt/atomgit/vllm.log",
            "detail": "Connection refused",
            "log": VLLM_LOG_PATH
        }

    except Timeout:
        # 服务在启动中或卡住
        return {
            "error": "vllm serve 响应超时，可能仍在加载模型，请稍后重试",
            "detail": "Request timeout",
            "log": VLLM_LOG_PATH
        }

    except Exception as e:
        # 兜底
        return {
            "error": "调用 vllm 时发生未知异常",
            "detail": str(e),
            "log": VLLM_LOG_PATH
        }

# base64数据
image_b64 = encode_file_to_base64(IMAGE_PATH)
video_b64 = encode_file_to_base64(VIDEO_PATH)
audio_b64 = encode_file_to_base64(AUDIO_PATH)

# 测试1: 纯文本对话
payload = {
    "messages": [
        {
            "role": "system",
            "content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
        },
        {
            "role": "user",
            "content": "who are you?"
        }
    ]
}
result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 读取本地图片并转为 Base64
def image_to_base64(image_path):
    with open(image_path, 'rb') as f:
        return base64.b64encode(f.read()).decode('utf-8')

# 替换为你的图片路径
image_path = 'example.jpg'
image_b64 = image_to_base64(image_path)
print(f"图片 Base64 长度: {len(image_b64)}")

payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What's in this image?"},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{image_b64}"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

payload = {
    "model": "Qwen2.5-Omni-7B",
    "messages": [
        {
            "role": "system",
            "content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "图片中是什么？"},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": "https://bpic.588ku.com/audio_copy/cover_img/18/11/04/0910b3378fa47cc657b69237c8573e42.png"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 测试4: 音频输入 (Base64编码, wav格式)
payload = {
    "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "input_audio",
                    "input_audio": {
                        "data": audio_b64,
                        "format": "wav"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 测试5: 音频URL
payload = {
    "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "描述一下音频"
                },
                {
                    "type": "audio_url",
                    "audio_url": {
                        "url": "https://isv-data.oss-cn-hangzhou.aliyuncs.com/ics/MaaS/ASR/test_audio/asr_example_zh.wav"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 测试6: 视频输入 (Base64编码)
payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "视频里面都有什么？"
                },
                {
                    "type": "video_url",
                    "video_url": {
                        "url": f"data:video/mp4;base64,{video_b64}"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 测试7: 视频URL
payload = {
    "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "描述一下视频"
                },
                {
                    "type": "video_url",
                    "video_url": {
                        "url": "https://vod.300hu.com/24/4c1f7a6atransbjngwcloud1oss/24ac9cd11020352573738127361/1097_4076_1_42660520d_f.mp4"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# content 支持同时传入多种媒体类型
"content": [
    {"type": "text", "text": "问题（涉及所有媒体）"},
    {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}},
    {"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}},
    {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}}
]

# 测试8: 混合媒体（图片+视频+音频）
payload = {
    "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
                }
            ]
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "图片里面有什么？音频里面有什么？视频里面有什么？"
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{image_b64}"
                    }
                },
                {
                    "type": "video_url",
                    "video_url": {
                        "url": f"data:video/mp4;base64,{video_b64}"
                    }
                },
                {
                    "type": "input_audio",
                    "input_audio": {
                        "data": audio_b64,
                        "format": "wav"
                    }
                }
            ]
        }
    ]
}

result = call_api(payload)
print(json.dumps(result, indent=2, ensure_ascii=False))

# 纯文本
payload = {
    "messages": [
        {"role": "user", "content": "问题"}
    ]
}

# 图片 + 文本
payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "问题"},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}}
            ]
        }
    ]
}

# 音频 + 文本
payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "问题"},
                {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}}
            ]
        }
    ]
}

# 视频 + 文本
payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "问题"},
                {"type": "video_url", "video_url": {"url": f"data:video/mp4;base64,{video_b64}"}}
            ]
        }
    ]
}

# 全模态混合（最强形态）
payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "综合描述图片、音频、视频内容"},
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
                {"type": "video_url", "video_url": {"url": f"data:video/mp4;base64,{video_b64}"}},
                {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}}
            ]
        }
    ]
}

# 启动后等待模型加载
import time
print("⏳ 等待模型加载...")
time.sleep(60)  # 至少等待 1 分钟
# 然后再发请求
result = call_api(payload)

# 1. 检查端口占用
!lsof -i :18004

# 2. 查看 vLLM 日志
!cat /opt/atomgit/vllm.log | tail -50

# 3. 杀掉旧进程
!kill -9 <PID>

# 1. 安装 atomgit
pip install -U atomgit -i https://mirrors.huaweicloud.com/repository/pypi/simple

# 2. 重启内核（执行后需重新运行所有 cell）
get_ipython().kernel.do_shutdown(restart=True)

# 3. 下载模型
from atomgit_hub import snapshot_download
snapshot_download(
    "hf_mirrors/Qwen/Qwen2.5-Omni-7B",
    local_dir="/opt/atomgit/Qwen2.5-Omni-7B"
)

# 4. 启动服务（防护脚本，一行命令）
PORT=18004 MODEL_DIR=/opt/atomgit/Qwen2.5-Omni-7B nohup vllm serve "hf://Qwen/Qwen2.5-Omni-7B" \
  --port 18004 \
  --tensor-parallel-size 1 \
  --dtype bfloat16 \
  --compilation-config '{"custom_ops":["none","+rms_norm","+rotary_embedding"]}' \
  --max-num-seqs 4 \
  --max-model-len 16000 \
  --gpu-memory-utilization 0.8 &

# 5. 下载测试数据
curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/image.jpg -o /opt/atomgit/static/image.jpg
curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/audio.wav -o /opt/atomgit/static/audio.wav
curl -L https://ai.gitcode.com/atomgit-ascend/qwen2.5-omni-7b/resolve/main/t-unit/video.mp4 -o /opt/atomgit/static/video.mp4