Update app.py

app.py

@@ -124,30 +124,46 @@ class ModelManager:
                 logger.info("Loading model with memory optimizations...")
                 clear_gpu_memory()
                 
-                # 모델 컴포넌트 로드 (메모리 효율적)
-                with torch.cuda.amp.autocast(enabled=False):
+                # 모델 컴포넌트 로드 (메모리 효율적) - autocast 수정
+                if torch.cuda.is_available():
+                    with torch.amp.autocast('cuda', enabled=False):  # 수정된 부분
+                        image_encoder = CLIPVisionModel.from_pretrained(
+                            config.model_id, 
+                            subfolder="image_encoder", 
+                            torch_dtype=torch.float16,
+                            low_cpu_mem_usage=True
+                        )
+                        
+                        vae = AutoencoderKLWan.from_pretrained(
+                            config.model_id, 
+                            subfolder="vae", 
+                            torch_dtype=torch.float16,
+                            low_cpu_mem_usage=True
+                        )
+                else:
+                    # CPU 환경
                     image_encoder = CLIPVisionModel.from_pretrained(
                         config.model_id, 
                         subfolder="image_encoder", 
-                        torch_dtype=torch.float16,  # float32 대신 float16 사용
+                        torch_dtype=torch.float32,
                         low_cpu_mem_usage=True
                     )
                     
                     vae = AutoencoderKLWan.from_pretrained(
                         config.model_id, 
                         subfolder="vae", 
-                        torch_dtype=torch.float16,  # float32 대신 float16 사용
+                        torch_dtype=torch.float32,
                         low_cpu_mem_usage=True
                     )
                     
-                    self._pipe = WanImageToVideoPipeline.from_pretrained(
-                        config.model_id, 
-                        vae=vae, 
-                        image_encoder=image_encoder, 
-                        torch_dtype=torch.bfloat16,
-                        low_cpu_mem_usage=True,
-                        use_safetensors=True
-                    )
+                self._pipe = WanImageToVideoPipeline.from_pretrained(
+                    config.model_id, 
+                    vae=vae, 
+                    image_encoder=image_encoder, 
+                    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+                    low_cpu_mem_usage=True,
+                    use_safetensors=True
+                )
                 
                 # 스케줄러 설정
                 self._pipe.scheduler = UniPCMultistepScheduler.from_config(
@@ -155,20 +171,24 @@ class ModelManager:
                 )
                 
                 # LoRA 로드
-                causvid_path = hf_hub_download(
-                    repo_id=config.lora_repo_id, filename=config.lora_filename
-                )
-                self._pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
-                self._pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
-                self._pipe.fuse_lora()
+                try:
+                    causvid_path = hf_hub_download(
+                        repo_id=config.lora_repo_id, filename=config.lora_filename
+                    )
+                    self._pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
+                    self._pipe.set_adapters(["causvid_lora"], adapter_weights=[0.95])
+                    self._pipe.fuse_lora()
+                    logger.info("LoRA weights loaded successfully")
+                except Exception as e:
+                    logger.warning(f"Failed to load LoRA weights: {e}")
                 
                 # GPU 최적화 설정
                 if hasattr(spaces, 'GPU'):  # Zero GPU 환경
                     self._pipe.enable_model_cpu_offload()
                     logger.info("CPU offload enabled for Zero GPU")
-                elif config.enable_model_cpu_offload:
+                elif config.enable_model_cpu_offload and torch.cuda.is_available():
                     self._pipe.enable_model_cpu_offload()
-                else:
+                elif torch.cuda.is_available():
                     self._pipe.to("cuda")
                 
                 if config.enable_vae_slicing:
@@ -284,9 +304,12 @@ class VideoGenerator:
                 required_memory = (height * width * 3 * 8 * duration * self.config.fixed_fps) / (1024**3)
                 if free_memory < required_memory * 2:
                     clear_gpu_memory()
-                    return False, "⚠️ Not enough GPU memory. Try smaller dimensions or shorter duration."
-            except:
-                pass  # GPU 체크 실패시 계속 진행
+                    # 재확인
+                    free_memory = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()
+                    if free_memory < required_memory * 1.5:
+                        return False, "⚠️ Not enough GPU memory. Try smaller dimensions or shorter duration."
+            except Exception as e:
+                logger.warning(f"GPU memory check failed: {e}")
         
         return True, None
     
@@ -394,26 +417,46 @@ def generate_video(input_image, prompt, height, width,
         progress(0.4, desc="🎬 Generating video frames...")
         
         # 메모리 효율적인 생성
-        with torch.inference_mode(), torch.cuda.amp.autocast(enabled=True):
-            try:
-                output_frames_list = pipe(
-                    image=resized_image,
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    height=target_h,
-                    width=target_w,
-                    num_frames=num_frames,
-                    guidance_scale=float(guidance_scale),
-                    num_inference_steps=int(steps),
-                    generator=torch.Generator(device="cuda").manual_seed(current_seed),
-                    return_dict=True
-                ).frames[0]
-            except torch.cuda.OutOfMemoryError:
-                clear_gpu_memory()
-                raise gr.Error("💾 GPU out of memory. Try smaller dimensions or shorter duration.")
-            except Exception as e:
-                logger.error(f"Generation error: {e}")
-                raise gr.Error(f"❌ Generation failed: {str(e)}")
+        if torch.cuda.is_available():
+            with torch.inference_mode(), torch.amp.autocast('cuda', enabled=True):  # 수정된 부분
+                try:
+                    output_frames_list = pipe(
+                        image=resized_image,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        height=target_h,
+                        width=target_w,
+                        num_frames=num_frames,
+                        guidance_scale=float(guidance_scale),
+                        num_inference_steps=int(steps),
+                        generator=torch.Generator(device="cuda").manual_seed(current_seed),
+                        return_dict=True
+                    ).frames[0]
+                except torch.cuda.OutOfMemoryError:
+                    clear_gpu_memory()
+                    raise gr.Error("💾 GPU out of memory. Try smaller dimensions or shorter duration.")
+                except Exception as e:
+                    logger.error(f"Generation error: {e}")
+                    raise gr.Error(f"❌ Generation failed: {str(e)}")
+        else:
+            # CPU 환경
+            with torch.inference_mode():
+                try:
+                    output_frames_list = pipe(
+                        image=resized_image,
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        height=target_h,
+                        width=target_w,
+                        num_frames=num_frames,
+                        guidance_scale=float(guidance_scale),
+                        num_inference_steps=int(steps),
+                        generator=torch.Generator().manual_seed(current_seed),
+                        return_dict=True
+                    ).frames[0]
+                except Exception as e:
+                    logger.error(f"Generation error: {e}")
+                    raise gr.Error(f"❌ Generation failed: {str(e)}")
         
         progress(0.9, desc="💾 Saving video...")
         filename = video_generator.generate_unique_filename(current_seed)
@@ -424,11 +467,19 @@ def generate_video(input_image, prompt, height, width,
         
         progress(1.0, desc="✨ Complete!")
         logger.info(f"Video generated successfully: {num_frames} frames, {target_h}x{target_w}")
+        
+        # 성공 정보 반환
+        info_text = f"✅ Generated {num_frames} frames at {target_h}x{target_w} with seed {current_seed}"
+        gr.Info(info_text)
+        
         return video_path, current_seed
         
+    except gr.Error:
+        # Gradio 에러는 그대로 전달
+        raise
     except Exception as e:
         logger.error(f"Unexpected error: {e}")
-        raise
+        raise gr.Error(f"❌ Unexpected error: {str(e)}")
         
     finally:
         # 항상 메모리 정리 및 락 해제
@@ -588,12 +639,34 @@ body {
     font-size: 0.9em;
 }
 
+.info-box {
+    background: rgba(52, 152, 219, 0.1);
+    border: 1px solid rgba(52, 152, 219, 0.3);
+    border-radius: 10px;
+    padding: 15px;
+    margin: 10px 0;
+    color: #2c5282;
+    font-size: 0.9em;
+}
+
 .footer {
     text-align: center;
     margin-top: 30px;
     color: #666;
     font-size: 0.9em;
 }
+
+/* 로딩 애니메이션 개선 */
+.progress-bar {
+    background: linear-gradient(90deg, #667eea 0%, #764ba2 50%, #667eea 100%);
+    background-size: 200% 100%;
+    animation: loading 1.5s ease-in-out infinite;
+}
+
+@keyframes loading {
+    0% { background-position: 0% 0%; }
+    100% { background-position: 200% 0%; }
+}
 """
 
 # Gradio UI
@@ -622,6 +695,19 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         </div>
         """)
         
+        # 새로운 정보 박스 추가
+        gr.HTML("""
+        <div class="info-box">
+            <strong>🎯 Quick Start Guide:</strong>
+            <ol style="margin: 5px 0; padding-left: 20px;">
+                <li>Upload your image - AI will calculate optimal dimensions</li>
+                <li>Enter a creative prompt or use the default</li>
+                <li>Adjust duration (1.5s recommended for best results)</li>
+                <li>Click Generate and wait ~60 seconds</li>
+            </ol>
+        </div>
+        """)
+        
         with gr.Row(elem_classes="main-content"):
             with gr.Column(scale=1):
                 gr.Markdown("### 📸 Input Settings")
@@ -718,21 +804,24 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                 
                 gr.HTML("""
                 <div class="footer">
-                    <p>💡 Tip: For best results, use clear images with good lighting</p>
+                    <p>💡 Tip: For best results, use clear images with good lighting and distinct subjects</p>
                 </div>
                 """)
         
-        # Examples
-        gr.Examples(
-            examples=[
-                ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 512, 512],
-                ["forg.jpg", "the frog jumps around", 576, 320],  # 16:9 aspect ratio within limits
-            ],
-            inputs=[input_image, prompt_input, height_slider, width_slider],
-            outputs=[video_output, seed],
-            fn=generate_video,
-            cache_examples=False  # 캐시 비활성화로 메모리 절약
-        )
+        # Examples - 파일명 확인 필요
+        try:
+            gr.Examples(
+                examples=[
+                    ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 512, 512],
+                    ["forg.jpg", "the frog jumps around", 576, 320],  # 16:9 aspect ratio within limits
+                ],
+                inputs=[input_image, prompt_input, height_slider, width_slider],
+                outputs=[video_output, seed],
+                fn=generate_video,
+                cache_examples=False  # 캐시 비활성화로 메모리 절약
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load examples: {e}")
 
         # 개선사항 요약 (작게)
         gr.HTML("""
@@ -768,4 +857,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()  # 더 엄격한 동시성 제어
+    # 초기 메모리 정리
+    clear_gpu_memory()
+    
+    # 앱 실행
+    demo.queue(concurrency_count=1)  # 동시 실행 제한
+    demo.launch()