Update app.py

app.py

@@ -4,13 +4,38 @@ import os
 import logging
 import json
 from datetime import datetime
+import tempfile
+import numpy as np
+from PIL import Image
 import shutil
+import httpx
+import time
+import base64
+from gradio_client import Client, handle_file
+import cv2
+from moviepy.editor import VideoFileClip
 
 # 로깅 설정
 logging.basicConfig(level=logging.INFO)
 
-# API 클라이언트 설정
-api_client = Client("http://211.233.58.202:7960/")
+# 타임아웃 설정을 30초로 늘림
+httpx_client = httpx.Client(timeout=30.0)
+
+max_retries = 3
+retry_delay = 5  # 5초 대기
+
+for attempt in range(max_retries):
+    try:
+        api_client = Client("http://211.233.58.202:7960/")
+        api_client.httpx_client = httpx_client  # httpx 클라이언트 설정
+        break  # 성공하면 루프 종료
+    except httpx.ReadTimeout:
+        if attempt < max_retries - 1:  # 마지막 시도가 아니면
+            print(f"Connection timed out. Retrying in {retry_delay} seconds...")
+            time.sleep(retry_delay)
+        else:
+            print("Failed to connect after multiple attempts.")
+            raise  # 모든 시도 실패 시 예외 발생
 
 # 갤러리 저장 디렉토리 설정
 GALLERY_DIR = "gallery"
@@ -19,16 +44,16 @@ GALLERY_JSON = "gallery.json"
 # 갤러리 디렉토리 생성
 os.makedirs(GALLERY_DIR, exist_ok=True)
 
-def save_to_gallery(image_path, prompt):
+def save_to_gallery(video_path, prompt):
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    new_image_path = os.path.join(GALLERY_DIR, f"{timestamp}.png")
+    new_video_path = os.path.join(GALLERY_DIR, f"{timestamp}.mp4")
     
-    # 이미지 파일 복사
-    shutil.copy2(image_path, new_image_path)
+    # 비디오 파일 복사
+    shutil.copy2(video_path, new_video_path)
     
     # 갤러리 정보 저장
     gallery_info = {
-        "image": new_image_path,
+        "video": new_video_path,
         "prompt": prompt,
         "timestamp": timestamp
     }
@@ -44,44 +69,55 @@ def save_to_gallery(image_path, prompt):
     with open(GALLERY_JSON, "w") as f:
         json.dump(gallery, f, indent=2)
     
-    return new_image_path
+    return new_video_path
 
 def load_gallery():
     if os.path.exists(GALLERY_JSON):
         with open(GALLERY_JSON, "r") as f:
             gallery = json.load(f)
-        return [(item["image"], item["prompt"]) for item in reversed(gallery)]
+        return [{"image": item["video"], "caption": item["prompt"]} for item in reversed(gallery)]
     return []
 
-def respond(message, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    logging.info(f"Received message: {message}, seed: {seed}, randomize_seed: {randomize_seed}, "
-                 f"width: {width}, height: {height}, guidance_scale: {guidance_scale}, "
-                 f"num_inference_steps: {num_inference_steps}")
+def update_gallery():
+    gallery_items = load_gallery()
+    return [
+        {"video": item['image'], "caption": item['caption']}
+        for item in gallery_items
+    ]
+
+def respond(image, prompt, steps, cfg_scale, eta, fs, seed, video_length):
+    logging.info(f"Received prompt: {prompt}, steps: {steps}, cfg_scale: {cfg_scale}, "
+                 f"eta: {eta}, fs: {fs}, seed: {seed}, video_length: {video_length}")
     
     try:
-        # 이미지 생성 요청
+        # 이미지 파일 처리
+        if image is not None:
+            image_file = handle_file(image)
+        else:
+            image_file = None
+
+        # 비디오 생성 요청
         result = api_client.predict(
-            prompt=message,
+            image=image_file,
+            prompt=prompt,
+            steps=steps,
+            cfg_scale=cfg_scale,
+            eta=eta,
+            fs=fs,
             seed=seed,
-            randomize_seed=randomize_seed,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            api_name="/infer_t2i"
+            video_length=video_length,
+            api_name="/infer"
         )
         logging.info("API response received: %s", result)
-        
-        # 결과 확인 및 처리
-        if isinstance(result, tuple) and len(result) >= 1:
-            image_path = result[0]
-            saved_image_path = save_to_gallery(image_path, message)
-            return saved_image_path
+
+        if isinstance(result, dict) and 'video' in result:
+            saved_video_path = save_to_gallery(result['video'], prompt)
+            return saved_video_path
         else:
-            raise ValueError("Unexpected API response format")
+            raise ValueError("예상치 못한 API 응답 형식")
     except Exception as e:
-        logging.error("Error during API request: %s", str(e))
-        return "Failed to generate image due to an error."
+        logging.error("API 요청 중 오류 발생: %s", str(e))
+        return "오류로 인해 비디오 생성에 실패했습니다."
 
 css = """
 footer {
@@ -119,64 +155,65 @@ examples = [
     ["A fantasy map of a fictional world, with detailed terrain and cities.", "q19.webp"]
 ]
 
-def use_prompt(prompt):
-    return prompt
+
+def use_prompt_and_image(prompt, image):
+    return prompt, image
+
+def show_video(evt: gr.SelectData):
+    return evt.value["video"] if isinstance(evt.value, dict) and "video" in evt.value else None
 
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
     with gr.Tab("Generate"):
         with gr.Row():
-            input_text = gr.Textbox(label="Enter your prompt for image generation")
-            output_image = gr.Image(label="Generated Image")
+            input_image = gr.Image(label="Upload an image", type="filepath")
+            input_text = gr.Textbox(label="Enter your prompt for video generation")
+            output_video = gr.Video(label="Generated Video")
         
         with gr.Row():
+            steps = gr.Slider(minimum=1, maximum=100, step=1, label="Steps", value=30)
+            cfg_scale = gr.Slider(minimum=1, maximum=15, step=0.1, label="CFG Scale", value=3.5)
+            eta = gr.Slider(minimum=0, maximum=1, step=0.1, label="ETA", value=1)
+            fs = gr.Slider(minimum=1, maximum=30, step=1, label="FPS", value=8)
             seed = gr.Slider(minimum=0, maximum=1000000, step=1, label="Seed", value=123)
-            randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
-        
-        with gr.Row():
-            width = gr.Slider(minimum=256, maximum=1024, step=64, label="Width", value=1024)
-            height = gr.Slider(minimum=256, maximum=1024, step=64, label="Height", value=576)
-        
-        with gr.Row():
-            guidance_scale = gr.Slider(minimum=1, maximum=20, step=0.1, label="Guidance Scale", value=5)
-            num_inference_steps = gr.Slider(minimum=1, maximum=100, step=1, label="Number of Inference Steps", value=28)
+            video_length = gr.Slider(minimum=1, maximum=10, step=1, label="Video Length (seconds)", value=2)
         
         with gr.Row():
             for prompt, image_file in examples:
                 with gr.Column():
                     gr.Image(image_file, label=prompt[:50] + "...")
-                    gr.Button("Use this prompt").click(
-                        fn=use_prompt,
+                    gr.Button("Use this example").click(
+                        fn=use_prompt_and_image,
                         inputs=[],
-                        outputs=input_text,
+                        outputs=[input_text, input_image],
                         api_name=False
                     ).then(
-                        lambda x=prompt: x,
+                        lambda p=prompt, i=image_file: (p, i),
                         inputs=[],
-                        outputs=input_text
+                        outputs=[input_text, input_image]
                     )
-    
+
     with gr.Tab("Gallery"):
         gallery = gr.Gallery(
-            label="Generated Images",
+            label="Generated Videos",
             show_label=False,
             elem_id="gallery",
             columns=[5],
             rows=[3],
             object_fit="contain",
-            height="auto"
+            height="auto",
+            preview=True
         )
+        selected_video = gr.Video(label="Selected Video")
         refresh_btn = gr.Button("Refresh Gallery")
 
-    def update_gallery():
-        return load_gallery()
-
     refresh_btn.click(fn=update_gallery, inputs=None, outputs=gallery)
     demo.load(fn=update_gallery, inputs=None, outputs=gallery)
+    gallery.select(show_video, None, selected_video)
 
     input_text.submit(
         fn=respond, 
-        inputs=[input_text, seed, randomize_seed, width, height, guidance_scale, num_inference_steps], 
-        outputs=output_image
+        inputs=[input_image, input_text, steps, cfg_scale, eta, fs, seed, video_length], 
+        outputs=output_video
     ).then(
         fn=update_gallery,
         inputs=None,