Upload app_hf_space_optimized.py

src/app_hf_space_optimized.py

@@ -0,0 +1,582 @@
+import streamlit as st
+import torch
+import gc
+import os
+import json
+import time
+import soundfile as sf
+from PIL import Image
+import numpy as np
+import ffmpeg # Use ffmpeg-python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from diffusers import StableDiffusionXLPipeline, CogVideoXPipeline
+from diffusers.utils import export_to_video
+from parler_tts import ParlerTTSForConditionalGeneration
+import tempfile
+import shutil
+import traceback
+import psutil # For memory stats
+
+st.set_page_config(layout="wide", page_title="POV Video Gen (HF Space)")
+
+# --- Configuration ---
+LLM_MODEL_ID = "Qwen/Qwen3-0.6B"
+IMAGE_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
+VIDEO_MODEL_ID = "THUDM/CogVideoX-2b"
+TTS_MODEL_ID = "parler-tts/parler-tts-mini-v1.1"
+
+IMAGE_WIDTH = 768
+IMAGE_HEIGHT = 1344
+SCENE_DURATION_SECONDS = 4 # Reduced duration for faster processing
+VIDEO_FPS = 10
+NUM_SCENES_DEFAULT = 3 # Lowered default
+MAX_SCENES = 4 # Stricter limit for free tier
+TEMP_SUBDIR = "pov_video_temp_hf" # Unique name
+
+# --- Device Setup & Memory Monitor ---
+mem_info_placeholder = st.sidebar.empty()
+
+def display_memory_usage():
+    """Displays CPU and GPU memory usage in the sidebar."""
+    try:
+        process = psutil.Process(os.getpid())
+        cpu_mem = process.memory_info().rss / (1024 * 1024) # MB
+        gpu_mem_info = "N/A"
+        if torch.cuda.is_available():
+            allocated = torch.cuda.memory_allocated(0) / (1024 * 1024) # MB
+            reserved = torch.cuda.memory_reserved(0) / (1024 * 1024) # MB
+            total = torch.cuda.get_device_properties(0).total_memory / (1024 * 1024) # MB
+            gpu_mem_info = f"Alloc: {allocated:.0f}MB | Reserv: {reserved:.0f}MB | Total: {total:.0f}MB"
+        mem_info_placeholder.info(f"🧠 CPU Mem: {cpu_mem:.0f} MB\n⚡ GPU Mem: {gpu_mem_info}")
+    except Exception as e:
+        mem_info_placeholder.warning(f"Could not get memory info: {e}")
+
+if torch.cuda.is_available():
+    device = "cuda"
+    try:
+        vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
+        st.sidebar.success(f"✅ GPU Detected! VRAM: {vram_gb:.2f} GB")
+        if vram_gb < 15:
+            st.sidebar.warning("⚠️ Low VRAM (< 15GB). May struggle.")
+    except Exception:
+        st.sidebar.warning("Could not read GPU VRAM.") # Continue assuming GPU exists
+else:
+    device = "cpu"
+    st.sidebar.error("⚠️ No GPU! App will be extremely slow & likely fail.")
+
+# --- Helper Functions ---
+def cleanup_gpu_memory(*args):
+    """Attempts to free GPU memory."""
+    print(f"Attempting GPU mem cleanup. Vars to del: {len(args)}")
+    display_memory_usage() # Before cleanup
+    del args # Remove reference to the tuple itself
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    display_memory_usage() # After cleanup
+    print("GPU mem cleanup done.")
+
+def get_temp_dir():
+    """Creates or returns the path to the temporary directory."""
+    # Use a consistent path within the app's execution context for simplicity on Spaces
+    # This might lead to leftover files if cleanup fails, but avoids potential permission issues with system temp
+    app_temp_dir = os.path.abspath(TEMP_SUBDIR) # Use relative path from script
+    os.makedirs(app_temp_dir, exist_ok=True)
+    if 'temp_dir_path' not in st.session_state or st.session_state.temp_dir_path != app_temp_dir:
+        print(f"Setting temp dir: {app_temp_dir}")
+        st.session_state.temp_dir_path = app_temp_dir
+    return app_temp_dir
+
+def cleanup_temp_dir():
+    """Removes the application's temporary directory."""
+    dir_path = st.session_state.get('temp_dir_path', None)
+    if dir_path and os.path.exists(dir_path) and TEMP_SUBDIR in dir_path: # Safety check
+        try:
+            shutil.rmtree(dir_path)
+            st.sidebar.success(f"Cleaned up: {dir_path}")
+            st.session_state.temp_dir_path = None
+        except Exception as e:
+            st.sidebar.error(f"Error cleaning temp dir {dir_path}: {e}")
+    else:
+         st.sidebar.info("Temp dir not found or already cleaned.")
+
+
+# --- Model Interaction Functions (Load -> Use -> Unload) ---
+
+def run_llm_step(user_prompt, num_scenes):
+    """Loads LLM, generates story, unloads LLM."""
+    st.info(f"🔄 Loading LLM: {LLM_MODEL_ID}...")
+    display_memory_usage()
+    llm_model, llm_tokenizer, model_inputs, generated_ids = None, None, None, None
+    story_data = None
+    try:
+        dtype = torch.bfloat16 if device=="cuda" and torch.cuda.is_bf16_supported() else torch.float16 if device=="cuda" else torch.float32
+        llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_ID)
+        llm_model = AutoModelForCausalLM.from_pretrained(
+            LLM_MODEL_ID, torch_dtype=dtype, low_cpu_mem_usage=True, device_map="auto" # Try low_cpu_mem_usage
+        )
+        display_memory_usage()
+        st.info("🧠 Generating story structure...")
+
+        # --- System Prompt --- (Shortened descriptions max length)
+        system_prompt = f"""
+You are an expert director creating POV TikTok video scripts.
+Break down the user's scenario into exactly {num_scenes} scenes ({SCENE_DURATION_SECONDS}s each).
+For EACH scene, generate:
+1. `scene_description`: Max 1-2 concise sentences describing action/setting for TTS. Max 350 characters.
+2. `image_prompt`: Detailed SDXL POV prompt (Start with "First-person perspective - pov shot of..."). Include setting, mood, style, time period, elements. Add "pov hands from the bottom corner..." if needed.
+3. `video_direction_prompt`: Simple camera action/motion for CogVideoX (e.g., "Camera pans right", "Subtle zoom in", "Static shot", "Hand reaches out").
+4. `audio_description`: Voice & ambience description for Parler-TTS (e.g., "Nervous male voice, faint market chatter.", "Calm female narrator, quiet library ambience.").
+
+Respond ONLY with a valid JSON object:
+{{
+"story_details": {{
+"title": "POV Title (Year)",
+"full_story": "Brief summary...",
+"scenes": [
+  {{ // Scene 1
+    "scene_description": "...", // Max 350 chars
+    "image_prompt": "...",
+    "video_direction_prompt": "...",
+    "audio_description": "..."
+  }},
+  // ... {num_scenes} scenes total ...
+]
+}}
+}}
+Strictly adhere to JSON format. No extra text.
+        """.strip()
+
+        messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": f"Create script: {user_prompt}"}]
+        text_input = llm_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        model_inputs = llm_tokenizer([text_input], return_tensors="pt").to(llm_model.device if hasattr(llm_model, 'device') else device)
+
+        # Use recommended parameters for non-thinking Qwen3
+        generated_ids = llm_model.generate(
+            **model_inputs, max_new_tokens=4096, # Still allow space for generation
+            temperature=0.7, top_p=0.8, top_k=20, do_sample=True,
+            pad_token_id=llm_tokenizer.eos_token_id # Important for stopping
+        )
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
+        response_text = llm_tokenizer.decode(output_ids, skip_special_tokens=True).strip()
+
+        st.write("LLM Raw Output:"); st.code(response_text, language='text')
+        json_string = response_text.strip().removeprefix("```json").removesuffix("```").strip()
+        parsed_data = json.loads(json_string)
+
+        if not ("story_details" in parsed_data and "scenes" in parsed_data["story_details"]): raise ValueError("Invalid JSON structure.")
+        actual_num_scenes = len(parsed_data["story_details"]["scenes"])
+        if actual_num_scenes != num_scenes: st.warning(f"LLM gave {actual_num_scenes} scenes, requested {num_scenes}.")
+
+        story_data = parsed_data["story_details"]
+        st.success("✅ Story generation complete.")
+    except Exception as e:
+        st.error(f"❌ LLM Step Failed: {e}"); st.error(traceback.format_exc()); story_data = None
+    finally:
+        st.info("🔄 Unloading LLM..."); cleanup_gpu_memory(llm_model, llm_tokenizer, model_inputs, generated_ids); st.info("✅ LLM Unloaded.")
+    return story_data
+
+def run_image_step(scenes, temp_dir):
+    st.info(f"🔄 Loading Image Generator: {IMAGE_MODEL_ID}...")
+    display_memory_usage()
+    image_pipe = None; image_results = []
+    try:
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        image_pipe = StableDiffusionXLPipeline.from_pretrained(
+            IMAGE_MODEL_ID, torch_dtype=dtype, use_safetensors=True, variant="fp16" if device == "cuda" else None,
+            low_cpu_mem_usage=True # Crucial for loading on low RAM systems
+        )
+        # Use CPU offloading even if it's slower, necessary for T4 VRAM
+        if device == "cuda": image_pipe.enable_model_cpu_offload()
+        else: image_pipe.to(device) # Move to CPU if needed
+        display_memory_usage()
+        st.info("🎨 Generating images sequentially...")
+
+        for i, scene in enumerate(scenes):
+            img_path = os.path.join(temp_dir, f"scene_{i+1}_img.png")
+            st.write(f"Generating Image {i+1}/{len(scenes)}...")
+            image = None # Define before try block
+            try:
+                with torch.no_grad():
+                    image = image_pipe(
+                        prompt=scene.get("image_prompt", "blank image"),
+                        width=IMAGE_WIDTH, height=IMAGE_HEIGHT, num_inference_steps=25 # Fewer steps for speed
+                    ).images[0]
+                image.save(img_path)
+                image_results.append({"scene": i, "path": img_path, "status": "succeeded"})
+                st.image(image, caption=f"Scene {i+1} OK", width=150)
+            except Exception as e:
+                st.error(f"❌ Image {i+1} Failed: {e}"); st.error(traceback.format_exc())
+                image_results.append({"scene": i, "path": None, "status": "failed"})
+            finally: cleanup_gpu_memory(image) # Clean intermediate var
+
+        st.success("✅ Image generation step complete.")
+    except Exception as e:
+        st.error(f"❌ Image Gen Step Failed: {e}"); st.error(traceback.format_exc())
+        image_results = [{"scene": i, "path": None, "status": "failed"} for i in range(len(scenes))]
+    finally:
+        st.info("🔄 Unloading Image Generator..."); cleanup_gpu_memory(image_pipe); st.info("✅ Image Generator Unloaded.")
+    return image_results
+
+def run_video_step(image_results, scenes, temp_dir):
+    successful_images = [item for item in image_results if item["status"] == "succeeded"]
+    if not successful_images: return []
+    st.info(f"🔄 Loading Video Generator: {VIDEO_MODEL_ID}...")
+    display_memory_usage()
+    video_pipe = None; video_results = []
+    try:
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        # Instantiate VAE and Transformer separately for potential offloading/quantization later if needed
+        # For now, load pipeline directly, enabling optimizations
+        video_pipe = CogVideoXPipeline.from_pretrained(VIDEO_MODEL_ID, torch_dtype=dtype)
+        if device == "cuda":
+             video_pipe.enable_model_cpu_offload()
+             video_pipe.enable_sequential_cpu_offload() # Needed for low VRAM
+        else: video_pipe.to(device)
+        video_pipe.vae.enable_slicing(); video_pipe.vae.enable_tiling()
+        display_memory_usage()
+        st.info("🎬 Generating videos sequentially...")
+        generator = torch.Generator(device=device)
+
+        for item in successful_images:
+            scene_index = item["scene"]; vid_path = os.path.join(temp_dir, f"scene_{scene_index + 1}_vid.mp4")
+            st.write(f"Generating Video for Scene {scene_index + 1}...")
+            img, video_frames = None, None # Define before try
+            try:
+                img = Image.open(item["path"])
+                video_direction = scenes[scene_index].get("video_direction_prompt", "subtle motion")
+                seed = int(time.time() * 1000 + scene_index) % 100000
+                if device == "cuda": generator.manual_seed(seed)
+                else: generator = torch.Generator(device='cpu').manual_seed(seed)
+
+                with torch.no_grad():
+                    video_frames = video_pipe(
+                        prompt=video_direction, image=img, num_inference_steps=40, # Slightly fewer steps
+                        num_frames=int(SCENE_DURATION_SECONDS * VIDEO_FPS) + 1,
+                        guidance_scale=6.0, generator=generator
+                    ).frames[0]
+                export_to_video(video_frames, vid_path, fps=VIDEO_FPS)
+                video_results.append({"scene": scene_index, "path": vid_path, "status": "succeeded"})
+                # Comment out preview to save resources on Spaces
+                # st.video(vid_path)
+                st.success(f"Video Scene {scene_index + 1} OK.")
+            except Exception as e:
+                st.error(f"❌ Video {scene_index + 1} Failed: {e}"); st.error(traceback.format_exc())
+                video_results.append({"scene": scene_index, "path": None, "status": "failed"})
+            finally: cleanup_gpu_memory(img, video_frames)
+
+        st.success("✅ Video generation step complete.")
+    except Exception as e:
+        st.error(f"❌ Video Gen Step Failed: {e}"); st.error(traceback.format_exc())
+        video_results = [{"scene": item["scene"], "path": None, "status": "failed"} for item in successful_images]
+    finally:
+        st.info("🔄 Unloading Video Generator..."); cleanup_gpu_memory(video_pipe); st.info("✅ Video Generator Unloaded.")
+    return video_results
+
+def run_audio_step(scenes, temp_dir):
+    st.info(f"🔄 Loading TTS Model: {TTS_MODEL_ID}...")
+    display_memory_usage()
+    tts_model, tts_tokenizer, tts_desc_tokenizer = None, None, None
+    audio_results = []
+    try:
+        # Load TTS model (Parler requires specific class)
+        tts_model = ParlerTTSForConditionalGeneration.from_pretrained(TTS_MODEL_ID).to(device)
+        tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL_ID) # For text prompt
+        tts_desc_tokenizer = AutoTokenizer.from_pretrained(tts_model.config.text_encoder._name_or_path) # For description
+        display_memory_usage()
+        st.info("🔊 Generating audio sequentially...")
+
+        for i, scene in enumerate(scenes):
+            audio_path = os.path.join(temp_dir, f"scene_{i+1}_audio.wav")
+            st.write(f"Generating Audio {i+1}/{len(scenes)}...")
+            desc_input_ids, prompt_input_ids, generation, audio_arr = None, None, None, None # Define before try
+            try:
+                text_to_speak = scene.get("scene_description", "")[:350] # Enforce limit
+                voice_description = scene.get("audio_description", "A neutral speaker.")
+                if not text_to_speak:
+                    audio_results.append({"scene": i, "path": None, "status": "skipped"})
+                    continue
+
+                desc_input_ids = tts_desc_tokenizer(voice_description, return_tensors="pt").input_ids.to(device)
+                prompt_input_ids = tts_tokenizer(text_to_speak, return_tensors="pt").input_ids.to(device)
+
+                with torch.no_grad():
+                    generation = tts_model.generate(
+                        input_ids=desc_input_ids, prompt_input_ids=prompt_input_ids,
+                        do_sample=True, temperature=0.7 # Slightly higher temp for variety
+                    ).to(torch.float32)
+
+                audio_arr = generation.cpu().numpy().squeeze()
+                sampling_rate = tts_model.config.sampling_rate
+                sf.write(audio_path, audio_arr, sampling_rate)
+                audio_results.append({"scene": i, "path": audio_path, "status": "succeeded"})
+                st.audio(audio_path, format='audio/wav') # Preview audio
+            except Exception as e:
+                st.error(f"❌ Audio {i+1} Failed: {e}"); st.error(traceback.format_exc())
+                audio_results.append({"scene": i, "path": None, "status": "failed"})
+            finally: cleanup_gpu_memory(desc_input_ids, prompt_input_ids, generation, audio_arr)
+
+        st.success("✅ Audio generation step complete.")
+    except Exception as e:
+        st.error(f"❌ Audio Gen Step Failed: {e}"); st.error(traceback.format_exc())
+        audio_results = [{"scene": i, "path": None, "status": "failed"} for i in range(len(scenes))]
+    finally:
+        st.info("🔄 Unloading TTS Model..."); cleanup_gpu_memory(tts_model, tts_tokenizer, tts_desc_tokenizer); st.info("✅ TTS Model Unloaded.")
+    return audio_results
+
+def run_compose_step_ffmpeg(video_results, audio_results, temp_dir, title="final_pov_video"):
+    """Combines videos and audio using ffmpeg-python."""
+    st.info("🎞️ Composing final video using ffmpeg-python (CPU)...")
+    display_memory_usage()
+    final_video_path = None
+    long_video_path = os.path.join(temp_dir, "long_video_temp.mp4")
+    long_audio_path = os.path.join(temp_dir, "long_audio_temp.wav")
+    final_output_path = os.path.join(temp_dir, f"{title}.mp4")
+    concat_video_list_path = os.path.join(temp_dir, "ffmpeg_video_list.txt")
+    concat_audio_list_path = os.path.join(temp_dir, "ffmpeg_audio_list.txt")
+
+    try:
+        successful_videos = sorted([item for item in video_results if item["status"] == "succeeded"], key=lambda x: x["scene"])
+        successful_audio = sorted([item for item in audio_results if item["status"] == "succeeded"], key=lambda x: x["scene"])
+
+        # Align based on scene index for safety
+        paths_to_compose = []
+        audio_map = {item['scene']: item['path'] for item in successful_audio}
+        for video_item in successful_videos:
+            scene_idx = video_item['scene']
+            if scene_idx in audio_map:
+                paths_to_compose.append({'scene': scene_idx, 'video': video_item['path'], 'audio': audio_map[scene_idx]})
+
+        if not paths_to_compose:
+            st.error("❌ No matching video/audio pairs found.")
+            return None
+
+        st.write(f"Found {len(paths_to_compose)} matching scene(s) to compose.")
+
+        # 1. Create file lists for ffmpeg concat demuxer
+        with open(concat_video_list_path, "w") as f_vid, open(concat_audio_list_path, "w") as f_aud:
+            for item in paths_to_compose:
+                f_vid.write(f"file '{os.path.relpath(item['video'], temp_dir)}'\n") # Use relative paths within temp dir
+                f_aud.write(f"file '{os.path.relpath(item['audio'], temp_dir)}'\n")
+
+        # 2. Concatenate Audio Files
+        st.write("Concatenating audio...")
+        try:
+            (
+                ffmpeg
+                .input(concat_audio_list_path, format='concat', safe=0, fflags='+igndts') # Add flags
+                .output(long_audio_path, acodec='pcm_s16le') # Output intermediate WAV
+                .global_args('-hide_banner', '-loglevel', 'error') # Suppress verbose output
+                .run(overwrite_output=True, cmd='ffmpeg') # Specify cmd='ffmpeg' if needed
+            )
+            st.write("Audio concatenated.")
+        except ffmpeg.Error as e:
+            st.error("FFmpeg Audio Concat Error:")
+            st.code(e.stderr.decode() if e.stderr else str(e))
+            raise # Re-raise to stop the process
+
+        # 3. Concatenate Video Files
+        st.write("Concatenating videos...")
+        try:
+            (
+                ffmpeg
+                .input(concat_video_list_path, format='concat', safe=0, fflags='+igndts')
+                .output(long_video_path, c='copy') # Use stream copy for speed
+                .global_args('-hide_banner', '-loglevel', 'error')
+                .run(overwrite_output=True, cmd='ffmpeg')
+            )
+            st.write("Videos concatenated.")
+        except ffmpeg.Error as e:
+            st.error("FFmpeg Video Concat Error:")
+            st.code(e.stderr.decode() if e.stderr else str(e))
+            raise
+
+        # 4. Mux (Combine) Video and Audio
+        st.write("Muxing final video...")
+        try:
+            in_video = ffmpeg.input(long_video_path)
+            in_audio = ffmpeg.input(long_audio_path)
+            (
+                ffmpeg
+                .output(in_video, in_audio, final_output_path, vcodec='copy', acodec='aac', shortest=None, strict='experimental') # Use aac audio codec
+                .global_args('-hide_banner', '-loglevel', 'error')
+                .run(overwrite_output=True, cmd='ffmpeg')
+            )
+            final_video_path = final_output_path # Set the final path on success
+            st.success("✅ Final video composed!")
+
+        except ffmpeg.Error as e:
+            st.error("FFmpeg Muxing Error:")
+            st.code(e.stderr.decode() if e.stderr else str(e))
+            final_video_path = None # Ensure it's None on failure
+            raise
+
+
+    except Exception as e:
+        st.error(f"❌ Video Composition Step Failed: {e}")
+        st.error(traceback.format_exc())
+        final_video_path = None
+    finally:
+        # Clean up intermediate files and lists
+        st.write("Cleaning up intermediate composition files...")
+        for f_path in [long_video_path, long_audio_path, concat_video_list_path, concat_audio_list_path]:
+             if os.path.exists(f_path):
+                 try: os.remove(f_path)
+                 except Exception as e_clean: print(f"Error cleaning {f_path}: {e_clean}")
+        display_memory_usage() # Final memory check for this step
+    return final_video_path
+
+
+# --- Streamlit UI ---
+
+st.title("🎬 POV Video Gen (HF Space Optimized)")
+st.caption("Local Generation: Scenario -> Story -> Images -> Videos -> Audio -> Compose -> Download")
+
+# Initialize Session State
+def init_state():
+    keys_to_init = {
+        'generation_in_progress': False, 'current_step': "idle", 'story_data': None,
+        'image_results': [], 'video_results': [], 'audio_results': [],
+        'final_video_path': None, 'temp_dir_path': None,
+        'num_scenes': NUM_SCENES_DEFAULT
+    }
+    for key, default_value in keys_to_init.items():
+        if key not in st.session_state:
+            st.session_state[key] = default_value
+init_state()
+
+# --- Sidebar ---
+with st.sidebar:
+    st.header("⚙️ Config & Control")
+    user_prompt = st.text_area("1. Enter POV Scenario:", height=100, value="POV: You're Marco Polo negotiating trade routes in the Silk Road bazaar (1270)", key="user_prompt_input")
+    num_scenes_req = st.number_input(f"2. Target Scenes (Max {MAX_SCENES}):", min_value=1, max_value=MAX_SCENES, value=st.session_state.num_scenes, key="num_scenes_req_input")
+
+    start_disable = st.session_state.generation_in_progress or device == "cpu"
+    start_button = st.button("🚀 Start Generation", type="primary", disabled=start_disable)
+
+    if start_button:
+        init_state() # Reset state variables first
+        st.session_state.generation_in_progress = True
+        st.session_state.current_step = "story"
+        st.session_state.num_scenes = num_scenes_req # Use the requested number
+        cleanup_temp_dir() # Clean old files
+        get_temp_dir() # Ensure new temp dir exists for this run
+        st.experimental_rerun()
+
+    st.header("⚠️ Actions")
+    if st.button("🔁 Reset Workflow", disabled=st.session_state.generation_in_progress):
+        init_state()
+        cleanup_temp_dir() # Also clean files on reset
+        st.experimental_rerun()
+
+    if st.button("🧹 Clean Temp Files Only", help=f"Removes files in {st.session_state.get('temp_dir_path', 'N/A')}", disabled=st.session_state.generation_in_progress):
+        cleanup_temp_dir()
+        st.experimental_rerun() # Rerun to update button help text etc.
+
+# --- Main Area Logic & Progress ---
+st.divider()
+if device == "cpu":
+    st.error("🔴 GPU (CUDA) is required. Cannot run on CPU.")
+elif st.session_state.generation_in_progress:
+    st.subheader(f"🚀 Running Step: {st.session_state.current_step.upper()}")
+    progress_bar = st.progress(0)
+    steps = ["story", "image", "video", "audio", "compose", "done"]
+    try:
+        current_index = steps.index(st.session_state.current_step)
+        progress_bar.progress((current_index / (len(steps) - 1)) * 100)
+    except ValueError:
+        progress_bar.progress(0) # Should not happen
+
+    # Use placeholders for status updates within each step function
+    status_placeholder = st.empty()
+
+    # Wrap the step execution in a try block to catch errors and stop
+    try:
+        temp_dir = get_temp_dir() # Ensure temp_dir is set
+        current_step = st.session_state.current_step # Local copy
+
+        if current_step == "story":
+            with status_placeholder.container(): st.session_state.story_data = run_llm_step(user_prompt, st.session_state.num_scenes)
+            next_step = "image" if st.session_state.story_data else "error"
+
+        elif current_step == "image":
+            scenes = st.session_state.story_data.get('scenes', [])
+            with status_placeholder.container(): st.session_state.image_results = run_image_step(scenes, temp_dir)
+            next_step = "video" if any(r['status'] == 'succeeded' for r in st.session_state.image_results) else "error"
+
+        elif current_step == "video":
+            scenes = st.session_state.story_data.get('scenes', [])
+            with status_placeholder.container(): st.session_state.video_results = run_video_step(st.session_state.image_results, scenes, temp_dir)
+            next_step = "audio" if any(r['status'] == 'succeeded' for r in st.session_state.video_results) else "error"
+
+        elif current_step == "audio":
+            scenes = st.session_state.story_data.get('scenes', [])
+            with status_placeholder.container(): st.session_state.audio_results = run_audio_step(scenes, temp_dir)
+            next_step = "compose" if any(r['status'] == 'succeeded' for r in st.session_state.audio_results) else "error"
+
+        elif current_step == "compose":
+            title_base = "".join(filter(str.isalnum, st.session_state.story_data.get('title', 'pov'))).replace(" ", "_") if st.session_state.story_data else "pov_video"
+            with status_placeholder.container(): st.session_state.final_video_path = run_compose_step_ffmpeg(
+                st.session_state.video_results, st.session_state.audio_results, temp_dir, title=title_base)
+            next_step = "done" if st.session_state.final_video_path else "error"
+
+        else: # Should not be reached if logic is right
+            next_step = "error"
+
+        # Update state and rerun ONLY if the step succeeded
+        if next_step != "error":
+             st.session_state.current_step = next_step
+             if next_step == "done":
+                 st.session_state.generation_in_progress = False # Workflow finished successfully
+                 progress_bar.progress(100)
+             st.experimental_rerun()
+        else:
+             st.error(f"🛑 Workflow failed at step: {current_step}")
+             st.session_state.current_step = "error"
+             st.session_state.generation_in_progress = False
+
+    except Exception as e:
+         st.error(f"An unexpected error occurred during step {st.session_state.current_step}: {e}")
+         st.error(traceback.format_exc())
+         st.session_state.current_step = "error"
+         st.session_state.generation_in_progress = False
+
+
+# --- Display Final Output ---
+st.divider()
+st.header("✅ Final Video")
+if st.session_state.current_step == "done" and st.session_state.final_video_path:
+    final_video_path = st.session_state.final_video_path
+    if os.path.exists(final_video_path):
+        st.video(final_video_path)
+        try:
+            with open(final_video_path, "rb") as fp:
+                st.download_button(
+                    label="⬇️ Download Final Video (.mp4)",
+                    data=fp,
+                    file_name=os.path.basename(final_video_path),
+                    mime="video/mp4",
+                    key="final_video_download_btn"
+                )
+        except Exception as e:
+            st.error(f"Error reading final video for download: {e}")
+    else:
+        st.error(f"Final video file not found: {final_video_path}. It might have been cleaned up.")
+elif st.session_state.current_step == "error":
+    st.error("🛑 Workflow failed. Check logs above. Please Reset and try again.")
+elif st.session_state.generation_in_progress:
+    st.info(f"⏳ Workflow running... Current step: **{st.session_state.current_step.upper()}**")
+else:
+    st.info("👋 Ready to generate. Use the sidebar to start.")
+
+# Optional: Display intermediate results in an expander
+with st.expander("Show Intermediate File Details", expanded=False):
+     st.write("**Story Data:**"); st.json(st.session_state.story_data or {})
+     st.write("**Image Results:**"); st.json(st.session_state.image_results or [])
+     st.write("**Video Results:**"); st.json(st.session_state.video_results or [])
+     st.write("**Audio Results:**"); st.json(st.session_state.audio_results or [])
+     st.write("**Final Path:**", st.session_state.final_video_path or "Not generated")
+     st.write("**Temp Dir:**", st.session_state.get('temp_dir_path', "N/A"))
+
+# Final memory display
+display_memory_usage()