Update app.py

app.py

@@ -24,16 +24,16 @@ YUNWU_API_KEY = "sk-Vhxjwm4XXu5fKrAtRNbRZGdPbocDZjG7B9UsSUjAdOQLyMUA"
 # 2. Sora 专用 API Key
 SORA_API_KEY = "sk-heZhMAAncKvJybPfhfx6rbj6ek0CoImJxrGPeRaXqSRpQR2t"
 
-# 模型选项
+# 模型配置
+TEXT_MODEL = "gemini-3-pro-preview-thinking"
 MODEL_OPTIONS = ["sora-2-all", "veo_3_1-fast"]
-TEXT_MODEL = "gemini-3-pro-preview-thinking" # 用于商品分析
 
 # 输出目录
-OUTPUT_DIR = "Ecommerce_Batch_Output"
+OUTPUT_DIR = "Ecommerce_Vertical_Output"
 
 # 并发配置
-MAX_WORKERS = 10      # 线程池大小
-VIDEO_WORKERS = 2     # 视频生成并发数 (建议保守设置)
+MAX_WORKERS = 10
+VIDEO_WORKERS = 2  # 视频生成并发数
 
 # ================= 提示词模版 (竖屏电商专用) =================
 
@@ -85,7 +85,6 @@ def image_to_base64(image_path):
 def image_to_data_uri(image_path):
     if not image_path: return None
     b64 = image_to_base64(image_path)
-    # 根据文件扩展名判断 mime type，默认 png
     return f"data:image/png;base64,{b64}"
 
 def download_file(url):
@@ -114,13 +113,13 @@ def clear_output_dir():
     if os.path.exists(OUTPUT_DIR): shutil.rmtree(OUTPUT_DIR)
     os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# ================= 核心 API 交互类 (双模型支持) =================
+# ================= 核心 API 交互类 =================
 
 class EcommerceDirector:
     def __init__(self, base_url):
         self.base_url = base_url
 
-    # Step 1: 分析图片和文本，生成 Prompts
+    # Step 1: 分析 (Gemini)
     def analyze_and_plan(self, image_path, description, count):
         headers = {
             "Authorization": f"Bearer {YUNWU_API_KEY}",
@@ -155,15 +154,18 @@ class EcommerceDirector:
             print(f"API Error: {e}")
             return [f"Showcase video of product {i+1} vertical style" for i in range(count)]
 
-    # Step 2: 视频生成路由器 (Router)
-    def generate_video(self, model_name, prompt, ref_image_path=None, use_ref=False):
+    # Step 2: 路由
+    def generate_video(self, model_name, prompt, ref_image_path):
+        if not ref_image_path:
+            return None, "Error: Reference image is mandatory."
+
         if "sora" in model_name.lower():
-            return self._generate_sora(model_name, prompt, ref_image_path, use_ref)
+            return self._generate_sora(model_name, prompt, ref_image_path)
         else:
-            return self._generate_veo(model_name, prompt, ref_image_path, use_ref)
+            return self._generate_veo(model_name, prompt, ref_image_path)
 
-    # === VEO 专用逻辑 (Multipart, 竖屏 9x16) ===
-    def _generate_veo(self, model_name, prompt, ref_image_path, use_ref):
+    # === VEO 逻辑 (9x16, Multipart) ===
+    def _generate_veo(self, model_name, prompt, ref_image_path):
         url = f"{self.base_url}/v1/videos"
         headers = {"Authorization": f"Bearer {YUNWU_API_KEY}"}
         
@@ -173,22 +175,17 @@ class EcommerceDirector:
                     'model': model_name,
                     'prompt': prompt,
                     'seconds': '5',
-                    'size': '9x16', # 【修改】竖屏
+                    'size': '9x16', # 竖屏
                     'watermark': 'false'
                 }
-                files = None
-                f_img = None
                 
-                if use_ref and ref_image_path:
-                    f_img = open(ref_image_path, 'rb')
+                with open(ref_image_path, 'rb') as f_img:
                     files = [('input_reference', (os.path.basename(ref_image_path), f_img, 'image/png'))]
-
-                resp = requests.post(url, headers=headers, data=data, files=files, timeout=120)
-                if f_img: f_img.close()
+                    resp = requests.post(url, headers=headers, data=data, files=files, timeout=120)
 
                 if resp.status_code == 200:
                     task_id = resp.json().get('id')
-                    return self._poll_result(task_id, headers)
+                    return self._poll_veo(task_id) # Veo 使用标准轮询
                 
                 print(f"[Veo] Submit Fail ({attempt}): {resp.text}")
                 time.sleep(2)
@@ -197,8 +194,8 @@ class EcommerceDirector:
                 time.sleep(2)
         return None, "Veo Failed"
 
-    # === SORA 专用逻辑 (JSON + Base64, 严格参照文档) ===
-    def _generate_sora(self, model_name, prompt, ref_image_path, use_ref):
+    # === SORA 逻辑 (Portrait, JSON+Base64, 独立 Query) ===
+    def _generate_sora(self, model_name, prompt, ref_image_path):
         url = f"{self.base_url}/v1/video/create"
         headers = {
             "Authorization": f"Bearer {SORA_API_KEY}",
@@ -208,31 +205,25 @@ class EcommerceDirector:
         
         for attempt in range(1, 4):
             try:
-                # 【修改】严格参照 OpenAPI 规范和截图
+                data_uri = image_to_data_uri(ref_image_path)
+                
                 payload = {
-                    "model": model_name, # sora-2-all
-                    "orientation": "portrait", # 【修改】竖屏
+                    "model": model_name,
+                    "orientation": "portrait", # 竖屏
                     "prompt": prompt,
                     "size": "large",
                     "duration": 5,
-                    "watermark": False, # 根据示例传递 Boolean
-                    "images": [] # 初始化数组
+                    "watermark": False,
+                    "images": [data_uri] # 强制垫图
                 }
-                
-                # Sora 垫图逻辑：Data URI 放入数组
-                if use_ref and ref_image_path:
-                    data_uri = image_to_data_uri(ref_image_path)
-                    if data_uri:
-                        payload["images"] = [data_uri]
 
-                # 发送 JSON 请求
                 resp = requests.post(url, headers=headers, json=payload, timeout=120)
                 
                 if resp.status_code == 200:
                     resp_json = resp.json()
                     task_id = resp_json.get('id')
                     if task_id:
-                        return self._poll_result(task_id, headers)
+                        return self._poll_sora(task_id) # Sora 使用特殊轮询
                 
                 print(f"[Sora] Submit Fail ({attempt}): {resp.text}")
                 time.sleep(2)
@@ -241,22 +232,39 @@ class EcommerceDirector:
                 time.sleep(2)
         return None, "Sora Failed"
 
-    # 通用轮询逻辑
-    def _poll_result(self, task_id, headers):
-        poll_url = f"{self.base_url}/v1/videos/{task_id}"
-        
-        for _ in range(60): # 轮询 3 分钟
+    # --- Veo 轮询 (标准) ---
+    def _poll_veo(self, task_id):
+        url = f"{self.base_url}/v1/videos/{task_id}"
+        headers = {"Authorization": f"Bearer {YUNWU_API_KEY}"}
+        return self._do_poll(url, headers)
+
+    # --- Sora 轮询 (Query 参数) ---
+    def _poll_sora(self, task_id):
+        # 按照文档：/v1/video/query?id=task_id
+        url = f"{self.base_url}/v1/video/query" 
+        headers = {
+            "Authorization": f"Bearer {SORA_API_KEY}",
+            "Accept": "application/json"
+        }
+        # requests params 会自动拼接 ?id=...
+        return self._do_poll(url, headers, params={"id": task_id})
+
+    # --- 通用轮询器 ---
+    def _do_poll(self, url, headers, params=None):
+        for _ in range(60): # 3分钟
             time.sleep(3)
             try:
-                resp = requests.get(poll_url, headers=headers)
+                resp = requests.get(url, headers=headers, params=params)
                 if resp.status_code == 200:
                     data = resp.json()
                     status = data.get('status')
+                    
                     if status in ['succeeded', 'success', 'completed']:
-                        url = self._deep_find_url(data)
-                        if url: return url, "OK"
+                        # 深度查找 video_url
+                        final_url = self._deep_find_url(data)
+                        if final_url: return final_url, "OK"
                     elif status == 'failed':
-                        return None, f"Remote Fail: {data.get('error')}"
+                        return None, f"Remote Fail: {data.get('error') or 'Unknown'}"
             except: pass
         return None, "Timeout"
 
@@ -282,55 +290,48 @@ director = EcommerceDirector(BASE_URL)
 
 def run_analysis_step(image, desc, count):
     if not image and not desc:
-        return "⚠️ 请至少上传图片或填写描述", None, gr.update(visible=False)
+        return "⚠️ 请上传图片或填写描述", None, gr.update(visible=False)
     
     logger.log(f"🕵️ Analyzing Product... Target: {count} videos")
     prompts = director.analyze_and_plan(image, desc, count)
+    
+    # Dataframe: [ID, Prompt]
     df_data = [[i+1, p] for i, p in enumerate(prompts)]
     
-    logger.log(f"✅ Analysis Done. Generated {len(prompts)} prompts.")
+    logger.log(f"✅ Generated {len(prompts)} prompts.")
     return logger.log("Ready to generate."), df_data, gr.update(visible=True)
 
-def run_generation_step(image, prompt_data, mode_str, model_name):
-    # 【修改】修复 DataFrame 真值判断错误
-    # prompt_data 可能是 None, List, 或 Pandas DataFrame
+def run_generation_step(image, prompt_data, model_name):
+    # 解析 Dataframe
     data_list = []
-    
     if prompt_data is None:
-        return "⚠️ 无有效提示词", None, "Failed"
-        
+        return "⚠️ 无提示词", None, "Failed"
     if isinstance(prompt_data, list):
         data_list = prompt_data
-    elif hasattr(prompt_data, 'values'): # 判断是否为 DataFrame
-        if prompt_data.empty:
-             return "⚠️ 提示词列表为空", None, "Failed"
+    elif hasattr(prompt_data, 'values'):
+        if prompt_data.empty: return "⚠️ 提示词为空", None, "Failed"
         data_list = prompt_data.values.tolist()
     
-    if len(data_list) == 0:
-        return "⚠️ 提示词列表为空", None, "Failed"
+    if len(data_list) == 0: return "⚠️ 列表为空", None, "Failed"
+
+    if not image:
+        return logger.log("⚠️ Error: 必须提供垫图 (Reference Image)"), None, "Error"
 
     clear_output_dir()
-    logger.log(f"🎬 Batch Start. Model: {model_name} | Mode: {mode_str} | Count: {len(data_list)}")
-    
-    use_ref = ("Image" in mode_str)
-    img_path = image if image else None
+    logger.log(f"🎬 Batch Start. Model: {model_name} | Count: {len(data_list)}")
     
-    if use_ref and not img_path:
-        return logger.log("⚠️ Error: Image mode selected but no image uploaded."), None, "Error"
-
     futures = []
     video_executor = concurrent.futures.ThreadPoolExecutor(max_workers=VIDEO_WORKERS)
     
-    # 提交任务
     for row in data_list:
         idx = row[0]
         prompt = row[1]
         logger.log(f"➕ Queueing Video {idx} ({model_name})...")
-        futures.append(video_executor.submit(process_single_video, idx, prompt, img_path, use_ref, model_name))
+        futures.append(video_executor.submit(process_single_video, idx, prompt, image, model_name))
 
-    # 等待结果
     completed = 0
     total = len(futures)
+    
     for f in concurrent.futures.as_completed(futures):
         idx, status = f.result()
         if status == "OK":
@@ -342,11 +343,12 @@ def run_generation_step(image, prompt_data, mode_str, model_name):
     video_executor.shutdown(wait=True)
     logger.log("📦 Zipping videos...")
     zip_path = create_zip(OUTPUT_DIR, "Ecommerce_Videos")
-    return logger.log("🎉 All Done! Ready to download."), zip_path, f"Completed {completed}/{total}"
+    
+    return logger.log("🎉 All Done!"), zip_path, f"Completed {completed}/{total}"
 
-def process_single_video(idx, prompt, img_path, use_ref, model_name):
+def process_single_video(idx, prompt, img_path, model_name):
     try:
-        url, msg = director.generate_video(model_name, prompt, img_path, use_ref)
+        url, msg = director.generate_video(model_name, prompt, img_path)
         if url:
             vid_bytes = download_file(url)
             if vid_bytes:
@@ -358,67 +360,46 @@ def process_single_video(idx, prompt, img_path, use_ref, model_name):
         print(f"Worker Error: {e}")
     return idx, "Fail"
 
-# ================= UI 界面 =================
+# ================= UI 界面 (默认颜色) =================
 
-dark_css = """
-body, .gradio-container { background-color: #0b0f19 !important; color: #e5e7eb !important; }
-.sidebar { background-color: #111827 !important; border-right: 1px solid #374151; padding: 20px; }
-.primary-btn { background: linear-gradient(90deg, #3b82f6, #2563eb) !important; border:none; color:white; font-weight:bold; }
-.secondary-btn { background-color: #374151 !important; color: white !important; border: 1px solid #4b5563 !important; }
-textarea, input { background-color: #1f2937 !important; color: #fff !important; border: 1px solid #374151 !important; }
-"""
-
-with gr.Blocks(title="Ecommerce Video Batch Agent", css=dark_css) as demo:
+with gr.Blocks(title="Ecommerce Video Generator") as demo:
     
     gr.Markdown("## 🛍️ 电商竖屏视频批量生成 (Sora-2 & Veo)")
+    gr.Markdown("单图生视频模式：Step 1 分析并生成分镜脚本 -> Step 2 使用主图批量生成视频")
     
     with gr.Row():
-        # 左侧配置区
-        with gr.Column(scale=1, elem_classes="sidebar"):
-            gr.Markdown("### 1. 商品信息输入")
-            input_image = gr.Image(label="商品图片 (Product Image)", type="filepath", height=250)
-            input_desc = gr.Textbox(label="商品描述 (Description)", placeholder="输入商品详情页卖点文案...", lines=4)
+        # 左侧配置
+        with gr.Column(scale=1):
+            gr.Markdown("### 1. 商品信息 (必填)")
+            input_image = gr.Image(label="商品主图 (必须上传，用于垫图)", type="filepath", height=250)
+            input_desc = gr.Textbox(label="商品描述", placeholder="输入商品卖点...", lines=4)
             
             gr.Markdown("### 2. 生成配置")
-            count_slider = gr.Slider(minimum=1, maximum=100, value=5, step=1, label="生成视频数量")
-            
-            # 模型选择
-            model_dropdown = gr.Dropdown(
-                choices=MODEL_OPTIONS, 
-                value="sora-2-all", 
-                label="选择视频模型 (Model)"
-            )
-            
-            # 模式选择
-            mode_dropdown = gr.Dropdown(
-                choices=["Text Only (纯文生)", "Image + Text (垫图生成)"], 
-                value="Image + Text (垫图生成)", 
-                label="生成模式 (Mode)"
-            )
+            count_slider = gr.Slider(minimum=1, maximum=100, value=5, step=1, label="生成数量")
+            model_dropdown = gr.Dropdown(choices=MODEL_OPTIONS, value="sora-2-all", label="视频模型")
             
-            analyze_btn = gr.Button("🔍 1. 分析并生成脚本", elem_classes="primary-btn")
+            analyze_btn = gr.Button("🔍 1. 分析并生成脚本", variant="primary")
             
-        # 右侧操作区
+        # 右侧操作
         with gr.Column(scale=2):
-            gr.Markdown("### 3. 脚本确认与修改")
+            gr.Markdown("### 3. 脚本确认")
             prompt_dataframe = gr.Dataframe(
                 headers=["ID", "Prompt"],
                 datatype=["number", "str"],
                 col_count=(2, "fixed"),
                 interactive=True,
-                label="生成的视频提示词 (可在生成前修改)",
-                wrap=True,
-                value=[[1, "Waiting for analysis..."]]
+                label="生成的分镜提示词 (可修改)",
+                value=[[1, "等待分析..."]]
             )
             
-            generate_btn = gr.Button("🎬 2. 开始批量生成", elem_classes="primary-btn", visible=False)
+            generate_btn = gr.Button("🎬 2. 开始批量生成 (使用主图)", variant="primary", visible=False)
             
-            gr.Markdown("### 4. 运行日志与下载")
-            log_box = gr.TextArea(label="系统日志", lines=8, interactive=False)
-            status_box = gr.Textbox(label="最终状态", interactive=False)
-            download_zip = gr.File(label="📦 下载视频包 (Download ZIP)", interactive=False)
+            gr.Markdown("### 4. 结果")
+            log_box = gr.TextArea(label="日志", lines=8, interactive=False)
+            status_box = gr.Textbox(label="状态", interactive=False)
+            download_zip = gr.File(label="下载视频包")
 
-    # 交互逻辑绑定
+    # 逻辑绑定
     analyze_btn.click(
         fn=run_analysis_step,
         inputs=[input_image, input_desc, count_slider],
@@ -427,7 +408,7 @@ with gr.Blocks(title="Ecommerce Video Batch Agent", css=dark_css) as demo:
     
     generate_btn.click(
         fn=run_generation_step,
-        inputs=[input_image, prompt_dataframe, mode_dropdown, model_dropdown],
+        inputs=[input_image, prompt_dataframe, model_dropdown],
         outputs=[log_box, download_zip, status_box]
     )