Update app.py

app.py

@@ -42,9 +42,9 @@ os.environ["HF_HOME"] = cache_path
 # CUDA 설정
 torch.backends.cuda.matmul.allow_tf32 = True
 
-# 번역기 초기화
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
+# 번역기 초기화 부분 수정
+translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cuda" if torch.cuda.is_available() else "cpu")
 
 if not path.exists(cache_path):
     os.makedirs(cache_path, exist_ok=True)
@@ -252,22 +252,110 @@ def save_image(image):
         logger.error(f"Error in save_image: {str(e)}")
         return None
 
-
-
 css = """
-footer {
-    visibility: hidden;
+footer {display: none}
+.gradio-container {max-width: 1200px !important}
+#gallery {
+    margin: 20px auto;
+    padding: 20px;
+}
+#gallery img {
+    width: 300px !important;
+    height: 300px !important;
+    object-fit: cover;
+    border-radius: 8px;
+}
+.gallery-item {
+    margin: 0 !important;
+    padding: 5px !important;
+}
+#video_player {
+    margin: 20px auto;
+    max-width: 800px;
+}
+.title {
+    text-align: center;
+    font-size: 1.5em;
+    margin: 10px 0;
 }
 """
 
 
-# Gradio 인터페이스 생성
+def get_random_seed():
+    return torch.randint(0, 1000000, (1,)).item()
+
+
+def create_thumbnail_gallery():
+    # 0부터 9까지의 이미지 리스트 생성
+    return [
+        "image/0.jpg", "image/1.jpg", "image/2.jpg",
+        "image/3.jpg", "image/4.jpg", "image/5.jpg",
+        "image/6.jpg", "image/7.jpg", "image/8.jpg",
+        "image/9.jpg"
+    ]
+
+def check_image_files():
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    missing_files = []
+    
+    for i in range(10):  # 0부터 9까지 확인
+        image_path = os.path.join(current_dir, f"image/{i}.jpg")
+        video_path = os.path.join(current_dir, f"image/{i}.mp4")
+        if not os.path.exists(image_path):
+            missing_files.append(f"{i}.jpg")
+        if not os.path.exists(video_path):
+            missing_files.append(f"{i}.mp4")
+    
+    if missing_files:
+        logger.error(f"Missing files: {', '.join(missing_files)}")
+        return False
+    return True
+
+def load_gallery_images():
+    gallery_images = []
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    
+    try:
+        for i in range(10):  # 0부터 9까지 로드
+            image_path = os.path.join(current_dir, f"image/{i}.jpg")
+            if os.path.exists(image_path):
+                img = Image.open(image_path)
+                gallery_images.append(img)
+            else:
+                logger.warning(f"Image not found: {image_path}")
+    except Exception as e:
+        logger.error(f"Error loading gallery images: {str(e)}")
+    
+    return gallery_images
+    
+# UI 부분 수정
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
     gr.HTML('<div class="title">🎥 Dokdo✨ Digital Odyssey from Korea, Designing Original</div>')
     gr.HTML('<div class="title">😄 Enjoy the amazing free video creation and enhancement services!</div>')
-    
+
     with gr.Tabs():
-        with gr.Tab("Image Generation"):
+        # 첫 번째 탭: Example Gallery
+        with gr.Tab("Example Gallery"):
+            with gr.Row():
+                gallery = gr.Gallery(
+                    value=create_thumbnail_gallery(),
+                    columns=[5],  # 한 줄에 5개씩 표시
+                    rows=[2],     # 2줄로 표시
+                    height="auto",
+                    show_label=False,
+                    elem_id="gallery"
+                )
+            
+            with gr.Row():
+                video_player = gr.Video(
+                    label="Selected Video",
+                    elem_id="video_player",
+                    interactive=False,
+                    autoplay=True
+                )
+
+        # 두 번째 탭: Image Generation
+        with gr.Tab("Image Generation & Enhanced"):
             with gr.Row():
                 with gr.Column(scale=3):
                     img_prompt = gr.Textbox(
@@ -278,61 +366,20 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                     
                     with gr.Accordion("Advanced Settings", open=False):
                         with gr.Row():
-                            height = gr.Slider(
-                                label="Height",
-                                minimum=256,
-                                maximum=1152,
-                                step=64,
-                                value=1024
-                            )
-                            width = gr.Slider(
-                                label="Width",
-                                minimum=256,
-                                maximum=1152,
-                                step=64,
-                                value=1024
-                            )
-                        
+                            height = gr.Slider(label="Height", minimum=256, maximum=1152, step=64, value=1024)
+                            width = gr.Slider(label="Width", minimum=256, maximum=1152, step=64, value=1024)
                         with gr.Row():
-                            steps = gr.Slider(
-                                label="Inference Steps",
-                                minimum=6,
-                                maximum=25,
-                                step=1,
-                                value=8
-                            )
-                            scales = gr.Slider(
-                                label="Guidance Scale",
-                                minimum=0.0,
-                                maximum=5.0,
-                                step=0.1,
-                                value=3.5
-                            )
-                        
-                        def get_random_seed():
-                            return torch.randint(0, 1000000, (1,)).item()
-                        
-                        seed = gr.Number(
-                            label="Seed",
-                            value=get_random_seed(),
-                            precision=0
-                        )
-                        
+                            steps = gr.Slider(label="Inference Steps", minimum=6, maximum=25, step=1, value=8)
+                            scales = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=5.0, step=0.1, value=3.5)
+                        seed = gr.Number(label="Seed", value=get_random_seed(), precision=0)
                         randomize_seed = gr.Button("🎲 Randomize Seed", elem_classes=["generate-btn"])
                     
-                    generate_btn = gr.Button(
-                        "✨ Generate Image",
-                        elem_classes=["generate-btn"]
-                    )
+                    generate_btn = gr.Button("✨ Generate Image", elem_classes=["generate-btn"])
                     
                 with gr.Column(scale=4):
-                    img_output = gr.Image(
-                        label="Generated Image",
-                        type="pil",
-                        format="png"
-                    )
-                    
+                    img_output = gr.Image(label="Generated Image", type="pil", format="png")
 
+        # 세 번째 탭: Video Generation
         with gr.Tab("Amazing Video Generation"):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -341,25 +388,19 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                         placeholder="비디오 설명을 입력하세요... (한글 입력 가능)",
                         lines=3
                     )
-                    upload_image = gr.Image(
-                        type="filepath",
-                        label="Upload First Frame Image"
-                    )
-                    video_generate_btn = gr.Button(
-                        "🎬 Generate Video",
-                        elem_classes=["generate-btn"]
-                    )
+                    upload_image = gr.Image(type="filepath", label="Upload First Frame Image")
+                    video_generate_btn = gr.Button("🎬 Generate Video", elem_classes=["generate-btn"])
                     
                 with gr.Column(scale=4):
                     video_output = gr.Video(label="Generated Video")
-                    
+
     @spaces.GPU
     def process_and_save_image(height, width, steps, scales, prompt, seed):
         is_safe, translated_prompt = process_prompt(prompt)
         if not is_safe:
             gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
             return None
-          
+      
         with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16), timer("inference"):
             try:
                 generated_image = pipe(
@@ -371,7 +412,7 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                     width=int(width),
                     max_sequence_length=256
                 ).images[0]
-                    
+                
                 if not isinstance(generated_image, Image.Image):
                     generated_image = Image.fromarray(generated_image)
                     
@@ -386,43 +427,50 @@ with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
                 logger.error(f"Error in image generation: {str(e)}")
                 return None    
 
-                        
-
     def process_and_generate_video(image, prompt):
         is_safe, translated_prompt = process_prompt(prompt)
         if not is_safe:
             gr.Warning("부적절한 내용이 포함된 프롬프트입니다.")
             return None
         return generate_video(image, translated_prompt)
-
+    
     def update_seed():
         return get_random_seed()
-        
+
+
+    # 이벤트 핸들러 수정
+    def show_video(evt: gr.SelectData):
+        video_num = evt.index  # 0부터 시작하는 인덱스
+        video_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"image/{video_num}.mp4")
+        if os.path.exists(video_path):
+            return video_path
+        return None
+
+    # 이벤트 연결
+    gallery.select(fn=show_video, outputs=video_player)
+
+
     generate_btn.click(
         process_and_save_image,
         inputs=[height, width, steps, scales, img_prompt, seed],
         outputs=img_output
     )
-
     video_generate_btn.click(
         process_and_generate_video,
         inputs=[upload_image, video_prompt],
         outputs=video_output
     )
-    
-    randomize_seed.click(
-        update_seed,
-        outputs=[seed]
-    )
-    
-    generate_btn.click(
-        update_seed,
-        outputs=[seed]
-    )
+    randomize_seed.click(update_seed, outputs=[seed])
+    generate_btn.click(update_seed, outputs=[seed])
 
 if __name__ == "__main__":
+    # 이미지와 비디오 파일 존재 확인
+    if not check_image_files():
+        print("Error: Required image and video files (0.jpg through 9.jpg and 0.mp4 through 9.mp4) are missing!")
+        sys.exit(1)
+        
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=True
+        share=False
     )