Update app.py

app.py

@@ -16,6 +16,11 @@ import hashlib
 from dataclasses import dataclass
 from typing import Optional, Tuple
 from functools import wraps
+import threading
+import os
+
+# GPU 메모리 관리 설정
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
 
 # 로깅 설정
 logging.basicConfig(level=logging.INFO)
@@ -40,10 +45,17 @@ class VideoGenerationConfig:
     max_frames: int = 81
     default_prompt: str = "make this image come alive, cinematic motion, smooth animation"
     default_negative_prompt: str = "static, blurred, low quality, watermark, text"
+    # GPU 메모리 최적화 설정
+    enable_model_cpu_offload: bool = True
+    enable_vae_slicing: bool = True
+    enable_vae_tiling: bool = True
 
 config = VideoGenerationConfig()
 MAX_SEED = np.iinfo(np.int32).max
 
+# 글로벌 락 (동시 실행 방지)
+generation_lock = threading.Lock()
+
 # 성능 측정 데코레이터
 def measure_time(func):
     @wraps(func)
@@ -54,12 +66,37 @@ def measure_time(func):
         return result
     return wrapper
 
-# 모델 관리자
+# GPU 메모리 정리 함수
+def clear_gpu_memory():
+    """강력한 GPU 메모리 정리"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+        gc.collect()
+        
+        # GPU 메모리 상태 로깅
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+        logger.info(f"GPU Memory - Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB")
+
+# 모델 관리자 (싱글톤 패턴)
 class ModelManager:
+    _instance = None
+    _lock = threading.Lock()
+    
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+    
     def __init__(self):
-        self._pipe = None
-        self._is_loaded = False
-        
+        if not hasattr(self, '_initialized'):
+            self._pipe = None
+            self._is_loaded = False
+            self._initialized = True
+    
     @property
     def pipe(self):
         if not self._is_loaded:
@@ -68,30 +105,93 @@ class ModelManager:
     
     @measure_time
     def _load_model(self):
-        logger.info("Loading model...")
-        image_encoder = CLIPVisionModel.from_pretrained(
-            config.model_id, subfolder="image_encoder", torch_dtype=torch.float32
-        )
-        vae = AutoencoderKLWan.from_pretrained(
-            config.model_id, subfolder="vae", torch_dtype=torch.float32
-        )
-        self._pipe = WanImageToVideoPipeline.from_pretrained(
-            config.model_id, vae=vae, image_encoder=image_encoder, torch_dtype=torch.bfloat16
-        )
-        self._pipe.scheduler = UniPCMultistepScheduler.from_config(
-            self._pipe.scheduler.config, flow_shift=8.0
-        )
-        self._pipe.to("cuda")
-        
-        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()
-        self._is_loaded = True
-        logger.info("Model loaded successfully")
+        """메모리 효율적인 모델 로딩"""
+        with self._lock:
+            if self._is_loaded:
+                return
+                
+            try:
+                logger.info("Loading model with memory optimizations...")
+                clear_gpu_memory()
+                
+                # 모델 컴포넌트 로드 (메모리 효율적)
+                with torch.cuda.amp.autocast(enabled=False):
+                    image_encoder = CLIPVisionModel.from_pretrained(
+                        config.model_id, 
+                        subfolder="image_encoder", 
+                        torch_dtype=torch.float16,  # float32 대신 float16 사용
+                        low_cpu_mem_usage=True
+                    )
+                    
+                    vae = AutoencoderKLWan.from_pretrained(
+                        config.model_id, 
+                        subfolder="vae", 
+                        torch_dtype=torch.float16,  # float32 대신 float16 사용
+                        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.scheduler = UniPCMultistepScheduler.from_config(
+                    self._pipe.scheduler.config, flow_shift=8.0
+                )
+                
+                # 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()
+                
+                # GPU 최적화 설정
+                if config.enable_model_cpu_offload:
+                    self._pipe.enable_model_cpu_offload()
+                else:
+                    self._pipe.to("cuda")
+                
+                if config.enable_vae_slicing:
+                    self._pipe.enable_vae_slicing()
+                    
+                if config.enable_vae_tiling:
+                    self._pipe.enable_vae_tiling()
+                
+                # xFormers 메모리 효율적인 attention 활성화 (가능한 경우)
+                try:
+                    self._pipe.enable_xformers_memory_efficient_attention()
+                    logger.info("xFormers memory efficient attention enabled")
+                except:
+                    logger.info("xFormers not available, using default attention")
+                
+                self._is_loaded = True
+                logger.info("Model loaded successfully with optimizations")
+                clear_gpu_memory()
+                
+            except Exception as e:
+                logger.error(f"Error loading model: {e}")
+                self._is_loaded = False
+                clear_gpu_memory()
+                raise
+    
+    def unload_model(self):
+        """모델 언로드 및 메모리 해제"""
+        with self._lock:
+            if self._pipe is not None:
+                del self._pipe
+                self._pipe = None
+                self._is_loaded = False
+                clear_gpu_memory()
+                logger.info("Model unloaded and memory cleared")
 
+# 싱글톤 인스턴스
 model_manager = ModelManager()
 
 # 비디오 생성기 클래스
@@ -136,6 +236,14 @@ class VideoGenerator:
         if duration > self.config.max_frames / self.config.fixed_fps:
             return False, f"⏱️ Duration too long (max {self.config.max_frames/self.config.fixed_fps:.1f}s)"
         
+        # GPU 메모리 체크
+        if torch.cuda.is_available():
+            free_memory = torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()
+            required_memory = (height * width * 3 * 8 * duration * config.fixed_fps) / (1024**3)  # 대략적인 추정
+            if free_memory < required_memory * 2:  # 2배 여유 확보
+                clear_gpu_memory()
+                return False, "⚠️ Not enough GPU memory. Try smaller dimensions or shorter duration."
+        
         return True, None
     
     def generate_unique_filename(self, seed: int) -> str:
@@ -165,12 +273,19 @@ def handle_image_upload(image):
 
 def get_duration(input_image, prompt, height, width, negative_prompt, 
                 duration_seconds, guidance_scale, steps, seed, randomize_seed, progress):
-    if steps > 4 and duration_seconds > 2:
-        return 90
-    elif steps > 4 or duration_seconds > 2:
-        return 75
-    else:
-        return 60
+    # GPU 사용량에 따라 동적으로 duration 조정
+    base_duration = 60
+    if steps > 4:
+        base_duration += 15
+    if duration_seconds > 2:
+        base_duration += 15
+    
+    # 해상도에 따른 추가 시간
+    pixels = height * width
+    if pixels > 500000:
+        base_duration += 20
+    
+    return min(base_duration, 120)  # 최대 120초
 
 @spaces.GPU(duration=get_duration)
 @measure_time
@@ -180,16 +295,23 @@ def generate_video(input_image, prompt, height, width,
                    seed=42, randomize_seed=False, 
                    progress=gr.Progress(track_tqdm=True)):
     
-    progress(0.1, desc="🔍 Validating inputs...")
-    
-    # 입력 검증
-    is_valid, error_msg = video_generator.validate_inputs(
-        input_image, prompt, height, width, duration_seconds, steps
-    )
-    if not is_valid:
-        raise gr.Error(error_msg)
+    # 동시 실행 방지
+    if not generation_lock.acquire(blocking=False):
+        raise gr.Error("⏳ Another video is being generated. Please wait...")
     
     try:
+        progress(0.1, desc="🔍 Validating inputs...")
+        
+        # 입력 검증
+        is_valid, error_msg = video_generator.validate_inputs(
+            input_image, prompt, height, width, duration_seconds, steps
+        )
+        if not is_valid:
+            raise gr.Error(error_msg)
+        
+        # 메모리 정리
+        clear_gpu_memory()
+        
         progress(0.2, desc="🎯 Preparing image...")
         target_h = max(config.mod_value, (int(height) // config.mod_value) * config.mod_value)
         target_w = max(config.mod_value, (int(width) // config.mod_value) * config.mod_value)
@@ -197,24 +319,35 @@ def generate_video(input_image, prompt, height, width,
                             config.min_frames, config.max_frames)
         current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
         
+        # 이미지 리사이즈 (메모리 효율적)
         resized_image = input_image.resize((target_w, target_h), Image.Resampling.LANCZOS)
         
         progress(0.3, desc="🎨 Loading model...")
         pipe = model_manager.pipe
         
         progress(0.4, desc="🎬 Generating video frames...")
-        with torch.inference_mode():
-            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)
-            ).frames[0]
+        
+        # 메모리 효율적인 생성
+        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)}")
         
         progress(0.9, desc="💾 Saving video...")
         filename = video_generator.generate_unique_filename(current_seed)
@@ -226,14 +359,23 @@ def generate_video(input_image, prompt, height, width,
         progress(1.0, desc="✨ Complete!")
         return video_path, current_seed
         
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        raise
+        
     finally:
+        # 항상 메모리 정리 및 락 해제
+        generation_lock.release()
+        
         # 메모리 정리
         if 'output_frames_list' in locals():
             del output_frames_list
-        gc.collect()
-        torch.cuda.empty_cache()
+        if 'resized_image' in locals():
+            del resized_image
+        
+        clear_gpu_memory()
 
-# CSS 스타일
+# 개선된 CSS 스타일
 css = """
 .container {
     max-width: 1200px;
@@ -249,17 +391,49 @@ css = """
     border-radius: 20px;
     color: white;
     box-shadow: 0 10px 30px rgba(0,0,0,0.2);
+    position: relative;
+    overflow: hidden;
+}
+
+.header::before {
+    content: '';
+    position: absolute;
+    top: -50%;
+    left: -50%;
+    width: 200%;
+    height: 200%;
+    background: radial-gradient(circle, rgba(255,255,255,0.1) 0%, transparent 70%);
+    animation: pulse 4s ease-in-out infinite;
+}
+
+@keyframes pulse {
+    0%, 100% { transform: scale(1); opacity: 0.5; }
+    50% { transform: scale(1.1); opacity: 0.8; }
 }
 
 .header h1 {
     font-size: 3em;
     margin-bottom: 10px;
     text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
+    position: relative;
+    z-index: 1;
 }
 
 .header p {
     font-size: 1.2em;
     opacity: 0.95;
+    position: relative;
+    z-index: 1;
+}
+
+.gpu-status {
+    position: absolute;
+    top: 10px;
+    right: 10px;
+    background: rgba(0,0,0,0.3);
+    padding: 5px 15px;
+    border-radius: 20px;
+    font-size: 0.8em;
 }
 
 .main-content {
@@ -296,6 +470,10 @@ css = """
     box-shadow: 0 7px 20px rgba(102, 126, 234, 0.6);
 }
 
+.generate-btn:active {
+    transform: translateY(0);
+}
+
 .video-output {
     background: #f8f9fa;
     padding: 20px;
@@ -333,8 +511,14 @@ body {
     100% { background-position: 0% 50%; }
 }
 
-.gr-button-secondary {
-    background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
+.warning-box {
+    background: rgba(255, 193, 7, 0.1);
+    border: 1px solid rgba(255, 193, 7, 0.3);
+    border-radius: 10px;
+    padding: 15px;
+    margin: 10px 0;
+    color: #856404;
+    font-size: 0.9em;
 }
 
 .footer {
@@ -348,11 +532,24 @@ body {
 # Gradio UI
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     with gr.Column(elem_classes="container"):
-        # Header
+        # Header with GPU status
         gr.HTML("""
         <div class="header">
             <h1>🎬 AI Video Magic Studio</h1>
             <p>Transform your images into captivating videos with Wan 2.1 + CausVid LoRA</p>
+            <div class="gpu-status">🖥️ GPU Optimized</div>
+        </div>
+        """)
+        
+        # GPU 메모리 경고
+        gr.HTML("""
+        <div class="warning-box">
+            <strong>💡 Performance Tips:</strong>
+            <ul style="margin: 5px 0; padding-left: 20px;">
+                <li>Start with lower resolution (512x512) for testing</li>
+                <li>Keep duration under 2 seconds for stable generation</li>
+                <li>Use 4-8 steps for optimal speed/quality balance</li>
+            </ul>
         </div>
         """)
         
@@ -409,14 +606,14 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                             maximum=config.slider_max_h,
                             step=config.mod_value,
                             value=config.default_height,
-                            label="📏 Height"
+                            label="📏 Height (lower = more stable)"
                         )
                         width_slider = gr.Slider(
                             minimum=config.slider_min_w,
                             maximum=config.slider_max_w,
                             step=config.mod_value,
                             value=config.default_width,
-                            label="📐 Width"
+                            label="📐 Width (lower = more stable)"
                         )
                     
                     steps_slider = gr.Slider(
@@ -459,58 +656,24 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
         # Examples
         gr.Examples(
             examples=[
-                ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 896, 512],
-                ["forg.jpg", "the frog jumps around", 448, 832],
+                ["peng.png", "a penguin playfully dancing in the snow, Antarctica", 512, 512],
+                ["forg.jpg", "the frog jumps around", 448, 448],
             ],
             inputs=[input_image, prompt_input, height_slider, width_slider],
             outputs=[video_output, seed],
             fn=generate_video,
-            cache_examples="lazy"
+            cache_examples=False  # 캐시 비활성화로 메모리 절약
         )
 
-    # Examples 섹션 후에 추가
-    gr.HTML("""
-    <div class="improvements-container" style="background: rgba(255, 255, 255, 0.95); backdrop-filter: blur(10px); border-radius: 15px; padding: 20px; margin: 20px auto; max-width: 800px; box-shadow: 0 5px 20px rgba(0,0,0,0.1); font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;">
-        <div class="improvements-header" style="text-align: center; margin-bottom: 20px;">
-            <h3 style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent; font-size: 1.5em; margin: 0; font-weight: 700;">✨ Enhanced Features</h3>
-            <p style="color: #666; font-size: 0.9em; margin-top: 5px;">Optimized for performance, stability, and user experience</p>
-        </div>
-        
-        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 15px;">
-            <div style="background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 10px; padding: 15px;">
-                <span style="font-size: 1.5em; margin-bottom: 8px; display: block;">🛡️</span>
-                <div style="font-weight: 600; color: #333; font-size: 0.95em; margin-bottom: 5px;">Robust Error Handling</div>
-                <div style="font-size: 0.75em; color: #666; line-height: 1.4;">Advanced validation and recovery mechanisms</div>
-            </div>
-            
-            <div style="background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 10px; padding: 15px;">
-                <span style="font-size: 1.5em; margin-bottom: 8px; display: block;">⚡</span>
-                <div style="font-weight: 600; color: #333; font-size: 0.95em; margin-bottom: 5px;">Performance Optimized</div>
-                <div style="font-size: 0.75em; color: #666; line-height: 1.4;">Faster processing with smart resource management</div>
-            </div>
-            
-            <div style="background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 10px; padding: 15px;">
-                <span style="font-size: 1.5em; margin-bottom: 8px; display: block;">🎨</span>
-                <div style="font-weight: 600; color: #333; font-size: 0.95em; margin-bottom: 5px;">Modern UI/UX</div>
-                <div style="font-size: 0.75em; color: #666; line-height: 1.4;">Beautiful interface with smooth animations</div>
-            </div>
-            
-            <div style="background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 10px; padding: 15px;">
-                <span style="font-size: 1.5em; margin-bottom: 8px; display: block;">🔧</span>
-                <div style="font-weight: 600; color: #333; font-size: 0.95em; margin-bottom: 5px;">Clean Architecture</div>
-                <div style="font-size: 0.75em; color: #666; line-height: 1.4;">Modular design for easy maintenance</div>
-            </div>
-        </div>
-        
-        <div style="display: flex; flex-wrap: wrap; gap: 5px; margin-top: 15px; justify-content: center;">
-            <span style="background: rgba(102, 126, 234, 0.1); color: #667eea; padding: 3px 10px; border-radius: 20px; font-size: 0.7em; font-weight: 500;">PyTorch</span>
-            <span style="background: rgba(102, 126, 234, 0.1); color: #667eea; padding: 3px 10px; border-radius: 20px; font-size: 0.7em; font-weight: 500;">Diffusers</span>
-            <span style="background: rgba(102, 126, 234, 0.1); color: #667eea; padding: 3px 10px; border-radius: 20px; font-size: 0.7em; font-weight: 500;">Gradio</span>
-            <span style="background: rgba(102, 126, 234, 0.1); color: #667eea; padding: 3px 10px; border-radius: 20px; font-size: 0.7em; font-weight: 500;">CUDA Optimized</span>
-            <span style="background: rgba(102, 126, 234, 0.1); color: #667eea; padding: 3px 10px; border-radius: 20px; font-size: 0.7em; font-weight: 500;">LoRA Enhanced</span>
+        # 개선사항 요약 (작게)
+        gr.HTML("""
+        <div style="background: rgba(255,255,255,0.9); border-radius: 10px; padding: 15px; margin-top: 20px; font-size: 0.8em; text-align: center;">
+            <p style="margin: 0; color: #666;">
+                <strong style="color: #667eea;">Enhanced with:</strong> 
+                🛡️ GPU Crash Protection • ⚡ Memory Optimization • 🎨 Modern UI • 🔧 Clean Architecture
+            </p>
         </div>
-    </div>
-    """)
+        """)
     
     # Event handlers
     input_image.upload(
@@ -536,4 +699,4 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue(max_size=1).launch()  # 큐 크기 제한으로 메모리 관리