Update app.py

app.py

@@ -7,6 +7,7 @@ import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
 import json
+import subprocess  # FFmpeg実行用に追加
 
 os.environ['HF_HOME'] = os.path.abspath(os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download')))
 
@@ -109,6 +110,76 @@ else:
 models = {}
 cpu_fallback_mode = not GPU_AVAILABLE  # GPUが利用できない場合、CPU代替モードを使用
 
+# 最後のフレームを抽出する関数を追加
+def extract_last_frame(video_path, output_path):
+    """
+    FFmpegを使用して動画の最後のフレームを抽出します
+    
+    Args:
+        video_path (str): 対象動画のパス
+        output_path (str): 出力画像のパス
+        
+    Returns:
+        bool: 抽出に成功した場合はTrue、失敗した場合はFalse
+    """
+    try:
+        # 動画からフレーム数を取得するコマンド
+        probe_cmd = [
+            'ffprobe', 
+            '-v', 'error', 
+            '-select_streams', 'v:0', 
+            '-count_packets', 
+            '-show_entries', 'stream=nb_read_packets', 
+            '-of', 'csv=p=0', 
+            video_path
+        ]
+        
+        # コマンドを実行してフレーム数を取得
+        frame_count = int(subprocess.check_output(probe_cmd).decode('utf-8').strip())
+        
+        # 最後のフレームを抽出するコマンド（0から始まるので-1する）
+        extract_cmd = [
+            'ffmpeg',
+            '-y',  # 既存ファイルを上書き
+            '-i', video_path,
+            '-vf', f'select=eq(n\\,{frame_count-1})',
+            '-vframes', '1',
+            output_path
+        ]
+        
+        # コマンドを実行
+        subprocess.run(extract_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        
+        # ファイルが生成されたか確認
+        if os.path.exists(output_path):
+            print(f"最後のフレームを抽出しました: {output_path}")
+            return True
+        else:
+            print(f"最後のフレーム抽出に失敗しました")
+            return False
+    except Exception as e:
+        print(f"フレーム抽出中にエラーが発生しました: {e}")
+        # 別の方法を試す（単純にシーク）
+        try:
+            simple_extract_cmd = [
+                'ffmpeg',
+                '-y',
+                '-sseof', '-1',  # 最後のフレームにシーク
+                '-i', video_path,
+                '-update', '1',
+                '-q:v', '1',
+                output_path
+            ]
+            subprocess.run(simple_extract_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            
+            if os.path.exists(output_path):
+                print(f"代替方法で最後のフレームを抽出しました: {output_path}")
+                return True
+        except Exception as e2:
+            print(f"代替フレーム抽出方法でもエラーが発生しました: {e2}")
+        
+        return False
+
 # モデルロード関数を使用
 def load_models():
     global models, cpu_fallback_mode, GPU_INITIALIZED
@@ -419,6 +490,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
 
     job_id = generate_timestamp()
     last_output_filename = None
+    last_frame_filename = None
     history_pixels = None
     history_latents = None
     total_generated_latent_frames = 0
@@ -588,7 +660,14 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                     try:
                         output_filename = os.path.join(outputs_folder, f'{job_id}_final_{total_generated_latent_frames}.mp4')
                         save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=18)
-                        stream.output_queue.push(('file', output_filename))
+                        
+                        # 最後のフレームを抽出して保存
+                        last_frame_filename = os.path.join(outputs_folder, f'{job_id}_final_{total_generated_latent_frames}_last_frame.png')
+                        extract_success = extract_last_frame(output_filename, last_frame_filename)
+                        
+                        # 抽出成功したかどうかに応じたファイル名を送信
+                        last_frame_path = last_frame_filename if extract_success else None
+                        stream.output_queue.push(('file', (output_filename, last_frame_path)))
                     except Exception as e:
                         print(f"最終動画保存中にエラーが発生しました: {e}")
                 
@@ -611,7 +690,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                 traceback.print_exc()
                 # 完全に終了せずに次のイテレーションを試みる
                 if last_output_filename:
-                    stream.output_queue.push(('file', last_output_filename))
+                    stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                 continue
 
             if not high_vram and not cpu_fallback_mode:
@@ -724,7 +803,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                     # 既に生成された動画がある場合、最後に生成された動画を返す
                     if last_output_filename:
                         print(f"【デバッグ】部分的に生成された動画あり: {last_output_filename}、この動画を返します")
-                        stream.output_queue.push(('file', last_output_filename))
+                        stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                         error_msg = "ユーザーにより生成プロセスが中断されましたが、部分的な動画は生成されています"
                     else:
                         print("【デバッグ】部分的に生成された動画なし、中断メッセージを返します")
@@ -742,7 +821,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                 
                 # 既に生成された動画がある場合、最後に生成された動画を返す
                 if last_output_filename:
-                    stream.output_queue.push(('file', last_output_filename))
+                    stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                     
                     # エラーメッセージを作成
                     error_msg = f"サンプリングプロセス中にエラーが発生しましたが、部分的に生成された動画を返します: {e}"
@@ -767,7 +846,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                 traceback.print_exc()
                 
                 if last_output_filename:
-                    stream.output_queue.push(('file', last_output_filename))
+                    stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                 stream.output_queue.push(('error', error_msg))
                 stream.output_queue.push(('end', None))
                 return
@@ -787,7 +866,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                 print(error_msg)
                 
                 if last_output_filename:
-                    stream.output_queue.push(('file', last_output_filename))
+                    stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                 continue
 
             try:
@@ -816,18 +895,24 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
                 save_start_time = time.time()
                 save_bcthw_as_mp4(history_pixels, output_filename, fps=30, crf=18)
                 print(f"動画保存完了、所要時間: {time.time() - save_start_time:.2f}秒")
-
+                
+                # 最後のフレームを抽出して保存
+                last_frame_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}_last_frame.png')
+                extract_success = extract_last_frame(output_filename, last_frame_filename)
+                
                 print(f'デコード完了。現在の潜在変数形状 {real_history_latents.shape}; ピクセル形状 {history_pixels.shape}')
 
                 last_output_filename = output_filename
-                stream.output_queue.push(('file', output_filename))
+                # 抽出成功したらファイル名も送信、失敗した場合はNoneを送信
+                last_frame_path = last_frame_filename if extract_success else None
+                stream.output_queue.push(('file', (output_filename, last_frame_path)))
             except Exception as e:
                 print(f"動画のデコードまたは保存中にエラーが発生しました: {e}")
                 traceback.print_exc()
                 
                 # 既に生成された動画がある場合、最後に生成された動画を返す
                 if last_output_filename:
-                    stream.output_queue.push(('file', last_output_filename))
+                    stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
                 
                 # エラー情報を記録
                 error_msg = f"動画のデコードまたは保存中にエラーが発生しました: {e}"
@@ -861,7 +946,7 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
         # 既に生成された動画がある場合、最後に生成された動画を返す
         if last_output_filename:
             print(f"【デバッグ】外部例外処理: 生成済み部分動画を返す {last_output_filename}")
-            stream.output_queue.push(('file', last_output_filename))
+            stream.output_queue.push(('file', (last_output_filename, last_frame_filename)))
         else:
             print("【デバッグ】外部例外処理: 生成済み動画が見つかりません")
         
@@ -891,8 +976,8 @@ if IN_HF_SPACE and 'spaces' in globals():
         gpu_memory_preservation = 6
 
 
-        # UI状態の初期化
-        yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
+        # UI状態の初期化 - last_frame_image用の更新を追加
+        yield None, None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
 
         try:
             stream = AsyncStream()
@@ -901,7 +986,9 @@ if IN_HF_SPACE and 'spaces' in globals():
             async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache)
 
             output_filename = None
+            last_frame_filename = None
             prev_output_filename = None
+            prev_last_frame_filename = None
             error_message = None
 
             # ワーカーの出力を継続的にチェック
@@ -910,15 +997,17 @@ if IN_HF_SPACE and 'spaces' in globals():
                     flag, data = stream.output_queue.next()
 
                     if flag == 'file':
-                        output_filename = data
+                        # data形式が(動画パス, 最後のフレームパス)に変更
+                        output_filename, last_frame_filename = data
                         prev_output_filename = output_filename
-                        # ファイル成功時にエラー表示をクリア
-                        yield output_filename, gr.update(), gr.update(), '', gr.update(interactive=False), gr.update(interactive=True)
+                        prev_last_frame_filename = last_frame_filename
+                        # ファイル成功時に出力を更新
+                        yield output_filename, last_frame_filename, gr.update(), gr.update(), '', gr.update(interactive=False), gr.update(interactive=True)
 
                     if flag == 'progress':
                         preview, desc, html = data
                         # 進捗更新時にエラーメッセージを変更せず、停止ボタンがインタラクティブであることを確認
-                        yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
+                        yield gr.update(), gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
                     
                     if flag == 'error':
                         error_message = data
@@ -926,16 +1015,17 @@ if IN_HF_SPACE and 'spaces' in globals():
                         # 即時表示せず、end信号を待機
 
                     if flag == 'end':
-                        # 最後の動画ファイルがある場合、確実に返す
+                        # 最後の動画ファイルと最後のフレーム画像がある場合、確実に返す
                         if output_filename is None and prev_output_filename is not None:
                             output_filename = prev_output_filename
+                            last_frame_filename = prev_last_frame_filename
                         
                         # エラーメッセージがある場合、わかりやすいエラー表示を作成
                         if error_message:
-                            yield output_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield output_filename, last_frame_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         else:
                             # 成功時にエラー表示をしない
-                            yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
+                            yield output_filename, last_frame_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
                         break
                 except Exception as e:
                     print(f"出力処理中にエラーが発生しました: {e}")
@@ -944,11 +1034,11 @@ if IN_HF_SPACE and 'spaces' in globals():
                     if current_time - last_update_time > 60:  # 60秒間更新がない場合、処理がフリーズした可能性
                         print(f"処理がフリーズした可能性があります。{current_time - last_update_time:.1f}秒間更新がありません")
                         
-                        # 部分的に生成された動画がある場合、それを返す
+                        # 部分的に生成された動画と最後のフレーム画像がある場合、それを返す
                         if prev_output_filename:
-                            yield prev_output_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield prev_output_filename, prev_last_frame_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         else:
-                            yield None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield None, None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         break
                     
         except Exception as e:
@@ -956,7 +1046,7 @@ if IN_HF_SPACE and 'spaces' in globals():
             traceback.print_exc()
             error_msg = str(e)
             
-            yield None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+            yield None, None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
     
     process = process_with_gpu
 else:
@@ -971,8 +1061,8 @@ else:
         rs = 0.0
         gpu_memory_preservation = 6
 
-        # UI状態の初期化
-        yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
+        # UI状態の初期化 - last_frame_image用の更新を追加
+        yield None, None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
 
         try:
             stream = AsyncStream()
@@ -981,7 +1071,9 @@ else:
             async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache)
 
             output_filename = None
+            last_frame_filename = None
             prev_output_filename = None
+            prev_last_frame_filename = None
             error_message = None
 
             # ワーカーの出力を継続的にチェック
@@ -990,15 +1082,17 @@ else:
                     flag, data = stream.output_queue.next()
 
                     if flag == 'file':
-                        output_filename = data
+                        # data形式が(動画パス, 最後のフレームパス)に変更
+                        output_filename, last_frame_filename = data
                         prev_output_filename = output_filename
-                        # ファイル成功時にエラー表示をクリア
-                        yield output_filename, gr.update(), gr.update(), '', gr.update(interactive=False), gr.update(interactive=True)
+                        prev_last_frame_filename = last_frame_filename
+                        # ファイル成功時の出力更新
+                        yield output_filename, last_frame_filename, gr.update(), gr.update(), '', gr.update(interactive=False), gr.update(interactive=True)
 
                     if flag == 'progress':
                         preview, desc, html = data
-                        # 進捗更新時にエラーメッセージを変更せず、停止ボタンがインタラクティブであることを確認
-                        yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
+                        # 進捗更新時
+                        yield gr.update(), gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
                     
                     if flag == 'error':
                         error_message = data
@@ -1006,16 +1100,17 @@ else:
                         # 即時表示せず、end信号を待機
 
                     if flag == 'end':
-                        # 最後の動画ファイルがある場合、確実に返す
+                        # 最後の動画ファイルと最後のフレーム画像がある場合、確実に返す
                         if output_filename is None and prev_output_filename is not None:
                             output_filename = prev_output_filename
+                            last_frame_filename = prev_last_frame_filename
                         
                         # エラーメッセージがある場合、わかりやすいエラー表示を作成
                         if error_message:
-                            yield output_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield output_filename, last_frame_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         else:
                             # 成功時にエラー表示をしない
-                            yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
+                            yield output_filename, last_frame_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
                         break
                 except Exception as e:
                     print(f"出力処理中にエラーが発生しました: {e}")
@@ -1024,11 +1119,11 @@ else:
                     if current_time - last_update_time > 60:  # 60秒間更新がない場合、処理がフリーズした可能性
                         print(f"処理がフリーズした可能性があります。{current_time - last_update_time:.1f}秒間更新がありません")
                         
-                        # 部分的に生成された動画がある場合、それを返す
+                        # 部分的に生成された動画と最後のフレーム画像がある場合、それを返す
                         if prev_output_filename:
-                            yield prev_output_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield prev_output_filename, prev_last_frame_filename, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         else:
-                            yield None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+                            yield None, None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
                         break
                     
         except Exception as e:
@@ -1036,7 +1131,7 @@ else:
             traceback.print_exc()
             error_msg = str(e)
             
-            yield None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
+            yield None, None, gr.update(visible=False), gr.update(), gr.update(interactive=True), gr.update(interactive=False)
     
 
 def end_process():
@@ -1219,6 +1314,32 @@ def make_custom_css():
     .error {
         display: none !important;
     }
+    
+    /* 最後のフレーム表示エリアのスタイル */
+    .last-frame-container {
+        margin-top: 15px;
+        border: 1px solid #ddd;
+        border-radius: 8px;
+        overflow: hidden;
+    }
+    
+    /* 最後のフレームとビデオの並べ表示（大画面用） */
+    @media (min-width: 1024px) {
+        .media-container {
+            display: flex;
+            flex-direction: row;
+            gap: 15px;
+        }
+        
+        .video-container {
+            flex: 2;
+        }
+        
+        .last-frame-container {
+            flex: 1;
+            max-width: 320px;
+        }
+    }
     """
     
     # CSSを結合
@@ -1285,16 +1406,16 @@ with block:
                 
                 seed = gr.Number(
                     label="シード値 / Seed", 
-                    value=31337, 
+                    value=1234, 
                     precision=0
                 )
 
                 # タッチ操作を最適化するためにslider-containerクラスを追加
                 with gr.Group(elem_classes="slider-container"):
                     total_second_length = gr.Slider(
-                        label="動画の長さ（最大3秒） / Video Length (max 3 seconds)", 
-                        minimum=0.5, 
-                        maximum=3, 
+                        label="動画の長さ（最大3秒） / Video Length (max 1.5 seconds)", 
+                        minimum=1.0, 
+                        maximum=1.5, 
                         value=1, 
                         step=0.1
                     )
@@ -1309,16 +1430,27 @@ with block:
                 elem_classes="preview-container"
             )
             
-            # 動画結果コンテナ
-            result_video = gr.Video(
-                label="生成された動画 / Generated Video", 
-                autoplay=True, 
-                show_share_button=True,  # 共有ボタンを追加
-                height=512, 
-                loop=True,
-                elem_classes="video-container",
-                elem_id="result-video"
-            )
+            # media-containerクラスを追加して動画と最後のフレームを並べて表示（大画面用）
+            with gr.Group(elem_classes="media-container"):
+                # 動画結果コンテナ
+                result_video = gr.Video(
+                    label="生成された動画 / Generated Video", 
+                    autoplay=True, 
+                    show_share_button=True,
+                    height=512, 
+                    loop=True,
+                    elem_classes="video-container",
+                    elem_id="result-video"
+                )
+                
+                # 新規追加：最後のフレーム表示用のImage要素
+                last_frame_image = gr.Image(
+                    label="最後のフレーム / Last Frame", 
+                    height=320,
+                    show_share_button=True,
+                    elem_classes="last-frame-container",
+                    elem_id="last-frame-image"
+                )
             
             gr.HTML("<div ='sampling_note' class='note'>注意：逆順サンプリングのため、終了動作が開始動作より先に生成されます。開始動作が動画に表示されていない場合は、しばらくお待ちください。後で生成されます。</div>")
             gr.HTML("<div ='sampling_note' class='note'>Note that the ending actions will be generated before the starting actions due to the inverted sampling. If the starting action is not in the video, you just need to wait, and it will be generated later.</div>")
@@ -1328,16 +1460,16 @@ with block:
                 progress_desc = gr.Markdown('', elem_classes='no-generating-animation')
                 progress_bar = gr.HTML('', elem_classes='no-generating-animation')
             
-            # エラーメッセージエリア - カスタムエラーメッセージ形式をサポートするHTMLコンポーネントを使用
+            # エラーメッセージエリア
             error_message = gr.HTML('', elem_id='error-message', visible=True)
     
-    # 処理関数
+    # 処理関数の入出力を更新
     ips = [input_image, prompt, n_prompt, seed, total_second_length, use_teacache]
+    ops = [result_video, last_frame_image, preview_image, progress_desc, progress_bar, start_button, end_button]
     
     # 開始と終了ボタンのイベント
-    start_button.click(fn=process, inputs=ips, outputs=[result_video, preview_image, progress_desc, progress_bar, start_button, end_button])
+    start_button.click(fn=process, inputs=ips, outputs=ops)
     end_button.click(fn=end_process)
 
 
-block.launch() 
-
+block.launch()