app.py

import os
import gc
import torch
import streamlit as st
import tempfile
import json
import subprocess
import shutil
from datetime import datetime
from io import BytesIO
import random
from PIL import Image
import base64 # For decoding images

# --- Hugging Face Model Libraries (Local Models for LLM & TTS) ---
from transformers import AutoTokenizer, AutoModelForCausalLM
from parler_tts import ParlerTTSForConditionalGeneration

# --- Google Generative AI (Gemini API) ---
try:
    from google import generativeai as genai
    from google.generativeai import types as genai_types
    google_gemini_sdk_available = True
except ImportError:
    google_gemini_sdk_available = False
    st.error("google-generativeai library not found. Please install it: pip install google-generativeai")

# --- Config ---
st.set_page_config(layout="wide", page_title="⚡ POV Generator Pro (Gemini SDK Images)")

# --- Local Model IDs (LLM & TTS) ---
LLM_MODEL_ID = "openai-community/gpt2-medium"
TTS_MODEL_ID = "parler-tts/parler-tts-mini-v1.1"

# --- Gemini API Configuration (from Streamlit Secrets) ---
# Ensure this is set in your Streamlit secrets (e.g., secrets.toml)
# Example secrets.toml:
# GEMINI_API_KEY = "YOUR_GEMINI_API_KEY"
# GEMINI_IMAGE_MODEL_ID = "gemini-2.0-flash-preview-image-generation" # Your custom/preview model

GEMINI_API_KEY = st.secrets.get("GEMINI_API_KEY", os.environ.get("GEMINI_API_KEY"))
# This is the model name you provided.
# WARNING: This is not a standard public model name for image generation via this SDK.
# It might be a preview model you have access to, or it might not work as expected for image generation.
GEMINI_IMAGE_MODEL_ID = st.secrets.get("GEMINI_IMAGE_MODEL_ID", "gemini-2.0-flash-preview-image-generation")


# --- Cache and Temp Directory Setup ---
CACHE_DIR = os.path.join(tempfile.gettempdir(), "hf_cache_pov_generator")
os.makedirs(CACHE_DIR, exist_ok=True)
os.environ['HUGGINGFACE_HUB_CACHE'] = CACHE_DIR
os.environ['HF_HOME'] = CACHE_DIR
os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR

# --- Session State Initialization ---
if 'run_id' not in st.session_state:
    st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
# ... (keep other session state initializations) ...
if 'story_data' not in st.session_state: st.session_state.story_data = None
if 'pil_images' not in st.session_state: st.session_state.pil_images = None
if 'image_paths_for_video' not in st.session_state: st.session_state.image_paths_for_video = None
if 'audio_paths' not in st.session_state: st.session_state.audio_paths = None
if 'video_path' not in st.session_state: st.session_state.video_path = None
if 'temp_base_dir' not in st.session_state: st.session_state.temp_base_dir = None


# --- Utility Functions (Keep from previous version) ---
def get_session_temp_dir():
    if st.session_state.get('temp_base_dir') and os.path.exists(st.session_state.temp_base_dir):
        return st.session_state.temp_base_dir
    base_dir = os.path.join(tempfile.gettempdir(), f"pov_generator_run_{st.session_state.run_id}")
    os.makedirs(base_dir, exist_ok=True)
    st.session_state.temp_base_dir = base_dir
    return base_dir

def cleanup_temp_files(specific_dir=None):
    path_to_clean = specific_dir or st.session_state.get("temp_base_dir")
    if path_to_clean and os.path.exists(path_to_clean):
        try: shutil.rmtree(path_to_clean)
        except Exception as e: st.warning(f"Error cleaning temp dir {path_to_clean}: {e}")
        if specific_dir is None: st.session_state.temp_base_dir = None

def clear_torch_cache():
    gc.collect()
    if torch.cuda.is_available(): torch.cuda.empty_cache()

# --- Model Loading (Cached for LLM & TTS - local) ---
@st.cache_resource
def load_llm_model_and_tokenizer(model_id):
    # ... (same as before)
    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR)
    model = AutoModelForCausalLM.from_pretrained(
        model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto", cache_dir=CACHE_DIR
    )
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = model.config.eos_token_id
    return model, tokenizer

@st.cache_resource
def load_tts_model_and_tokenizers(model_id):
    # ... (same as before)
    tts_model = ParlerTTSForConditionalGeneration.from_pretrained(
        model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto", cache_dir=CACHE_DIR
    )
    prompt_tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR)
    desc_tokenizer_path = tts_model.config.text_encoder.name_or_path if hasattr(tts_model.config.text_encoder, 'name_or_path') else tts_model.config.text_encoder._name_or_path
    desc_tokenizer = AutoTokenizer.from_pretrained(desc_tokenizer_path, cache_dir=CACHE_DIR)
    return tts_model, prompt_tokenizer, desc_tokenizer

# --- Gemini SDK Client (Cached) ---
@st.cache_resource
def get_gemini_client(api_key):
    if not google_gemini_sdk_available: return None
    if not api_key:
        st.error("GEMINI_API_KEY not found in Streamlit secrets. Please configure it.")
        return None
    try:
        genai.configure(api_key=api_key)
        # The client itself isn't what we cache, genai.configure sets it globally for the module.
        # We just need a way to ensure it's configured once.
        # For actual model interaction, we'll use genai.GenerativeModel
        return genai # Return the module itself as a signal it's configured
    except Exception as e:
        st.error(f"Error configuring Gemini API client: {e}")
        return None

# --- Step 1: Generate JSON Story (Keep from previous version) ---
def generate_story(prompt: str, num_scenes: int):
    # ... (Same as your previous generate_story function)
    model, tokenizer = load_llm_model_and_tokenizer(LLM_MODEL_ID)
    story_prompt = (
        f"Generate a compelling short POV story based on the following prompt: '{prompt}'. "
        f"The story should consist of exactly {num_scenes} distinct scenes. "
        f"Clearly separate each scene with the delimiter '###'. "
        f"Do not include any introductory or concluding text. Each scene: 2-4 sentences."
    )
    input_ids = tokenizer.encode(story_prompt, return_tensors="pt").to(model.device)
    max_model_tokens = getattr(model.config, 'n_positions', 1024)
    max_new = max_model_tokens - input_ids.shape[1] - 20
    max_new_tokens_val = min(num_scenes * 75, 700, max_new)

    if max_new_tokens_val <= 0: st.error("Prompt too long."); return None
    output = model.generate(
        input_ids, max_new_tokens=max_new_tokens_val, do_sample=True,
        temperature=0.7, top_k=50, pad_token_id=tokenizer.eos_token_id
    )
    full_result = tokenizer.decode(output[0], skip_special_tokens=True)
    generated_text = full_result[len(story_prompt):].strip() if full_result.startswith(story_prompt) else full_result
    
    scenes_raw = generated_text.split("###")
    processed_scenes = [s.strip() for s in scenes_raw if s.strip()]
    
    if not processed_scenes: st.error(f"No scenes from LLM: {generated_text}"); return None
    final_scenes = processed_scenes[:num_scenes]
    if len(final_scenes) < num_scenes: st.warning(f"LLM made {len(final_scenes)}/{num_scenes} scenes.")
    
    clear_torch_cache()
    return {"title": prompt[:60].capitalize(), "scenes": final_scenes}


# --- Step 2: Generate Images with Gemini SDK ---
def generate_image_with_gemini_sdk(prompt_text: str, seed: int): # Width/Height might not be controllable
    if not google_gemini_sdk_available:
        st.error("Google Generative AI SDK not available.")
        return None
    
    gemini_configured_module = get_gemini_client(GEMINI_API_KEY)
    if not gemini_configured_module:
        return None

    try:
        # IMPORTANT: The model name GEMINI_IMAGE_MODEL_ID is user-provided and may not be a standard image generation model.
        # The `response_modalities` config is also based on the user's snippet.
        # This whole block assumes this model and config work as intended by the user.
        model_instance = genai.GenerativeModel(GEMINI_IMAGE_MODEL_ID)

        st.write(f"Using Gemini Model for Image: {GEMINI_IMAGE_MODEL_ID}") # For debugging
        # The concept of 'seed' might not be directly applicable or exposed in the same way
        # for `generate_content` with some Gemini models. We include it if the API might use it.
        # GenerationConfig can also take 'temperature', 'top_p', 'top_k', 'candidate_count', 'max_output_tokens'.
        # 'seed' is not a standard param for GenerateContentConfig in the public SDK.
        # We will use a basic config based on your snippet.
        
        generation_config = genai_types.GenerationConfig(
            # response_mime_type="image/png", # Often used with specific image models, but response_modalities is different
            # Not adding seed here as it's not standard for this config object.
            # Temperature, etc., could be added if relevant for image quality.
        )

        # The `contents` structure might need to be more complex if the model expects specific formatting.
        # For a simple prompt, a string might suffice.
        # The `response_modalities` part is from your snippet, which is unusual for GenerateContentConfig.
        # It's typically set at the model level or inferred.
        # Let's try to construct the request as closely as possible to your example structure.
        # `client.models.generate_content` is not the typical way; it's usually `GenerativeModel_instance.generate_content`
        
        # Corrected way to use generate_content with a GenerativeModel instance:
        response = model_instance.generate_content(
            contents=[prompt_text], # `contents` should be an iterable of Content objects or strings
            generation_config=generation_config,
            # The `response_modalities` is not a direct parameter of `generate_content`.
            # It's usually an expectation from the model. If your model *requires* it,
            # it might need to be part of a more complex `Content` object within `contents`.
            # For now, we'll assume the model inherently knows to produce an image if prompted correctly,
            # or that `GEMINI_IMAGE_MODEL_ID` implies image output.
        )

        # Process the response to find image data
        # A response can have multiple `parts`. Image data is typically in a part with a specific mime_type.
        if response.parts:
            for part in response.parts:
                if part.mime_type and part.mime_type.startswith("image/"):
                    # Image data is in `part.data` (bytes) or `part.inline_data.data`
                    image_bytes = None
                    if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'): # Common structure
                        image_bytes = part.inline_data.data
                    elif hasattr(part, 'data'): # Fallback
                        image_bytes = part.data
                    
                    if image_bytes:
                        pil_image = Image.open(BytesIO(image_bytes))
                        return pil_image
                    else:
                        st.warning("Image part found in Gemini response, but no image data.")
            st.error("No image data found in Gemini response parts, though parts existed.")
            st.json([part_to_dict(p) for p in response.parts]) # Show what was received
            return None
        else:
            st.error("Gemini SDK returned no parts in the response.")
            if response.prompt_feedback:
                st.warning(f"Prompt Feedback: {response.prompt_feedback}")
            return None

    except Exception as e:
        st.error(f"Error during Gemini SDK image generation: {e}")
        st.error(f"Model used: {GEMINI_IMAGE_MODEL_ID}. Prompt: {prompt_text[:100]}...")
        if "API key not valid" in str(e) or "PERMISSION_DENIED" in str(e):
            st.error("Please check your GEMINI_API_KEY and ensure it's valid and has permissions for this model.")
        elif "Could not find model" in str(e) or "MODEL_NAME_INVALID" in str(e):
            st.error(f"The model '{GEMINI_IMAGE_MODEL_ID}' could not be found or is invalid for this API. Please verify the model name.")
        elif "DeadlineExceeded" in str(e):
            st.error("The request to Gemini API timed out. The model might be slow or the task too complex.")
        # You might want to inspect the full error for more details
        # print(f"Full Gemini API Error: {e}")
        return None

def part_to_dict(part):
    """Helper to convert a Part object to a dict for st.json, handling bytes."""
    data = " данни" # Placeholder if data is binary and large
    if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
        data = f"<image bytes: {len(part.inline_data.data)}>" if part.inline_data.data else "няма данни"
    elif hasattr(part, 'data'):
        data = f"<bytes: {len(part.data)}>" if part.data else "няма данни"
    
    return {
        "text": part.text if hasattr(part, 'text') else None,
        "mime_type": part.mime_type if hasattr(part, 'mime_type') else None,
        "data_summary": data
    }


def generate_images_for_scenes(scenes, randomize_seed_img, base_seed_img): # img_width, img_height removed as likely not controllable
    pil_images = []
    frames_dir = os.path.join(get_session_temp_dir(), "frames_for_video")
    os.makedirs(frames_dir, exist_ok=True)
    image_paths_for_video = []

    cols = st.columns(min(3, len(scenes)))
    col_idx = 0

    st.warning(f"""
        **Note on Image Generation with `{GEMINI_IMAGE_MODEL_ID}`:**
        - The model name provided is not a standard public image generation model via this SDK. Results may vary.
        - Image dimensions (width/height) are determined by the model. UI sliders for width/height are not used for this generation method.
        - The 'seed' parameter's effect depends on the specific model's implementation via this API.
    """)

    for i, scene_text in enumerate(scenes):
        # Seed concept might be different or not directly exposed for `generate_content`
        current_seed = random.randint(0, 9_999_999) if randomize_seed_img else base_seed_img + i
        
        with st.spinner(f"Generating image for scene {i+1} via Gemini SDK (seed hint: {current_seed})..."):
            img = None
            try:
                # The prompt needs to clearly ask for an image based on the scene text.
                image_prompt = f"Visually represent this scene: {scene_text}. Cinematic, detailed, high quality."
                img = generate_image_with_gemini_sdk(image_prompt, current_seed)
            except Exception as e:
                st.error(f"Failed to generate image for scene {i+1}: {e}")
                img = None

            pil_images.append(img)
            if img:
                st.caption(f"Generated image size: {img.width}x{img.height}")
                img_path = os.path.join(frames_dir, f"frame_{i:03d}.png")
                img.save(img_path)
                image_paths_for_video.append(img_path)
                
                with cols[col_idx % len(cols)]:
                    st.image(img, caption=f"Scene {i+1} (Seed hint: {current_seed})")
                    img_byte_arr = BytesIO()
                    img.save(img_byte_arr, format='PNG')
                    st.download_button(
                        label=f"Download Img {i+1}", data=img_byte_arr.getvalue(),
                        file_name=f"scene_{i+1}_image.png", mime="image/png", key=f"dl_img_{i}"
                    )
                col_idx += 1
            else:
                image_paths_for_video.append(None)
                with cols[col_idx % len(cols)]:
                    st.warning(f"Image for Scene {i+1} could not be generated.")
                col_idx += 1
    return pil_images, image_paths_for_video

# --- Step 3: Generate TTS (Keep from previous version) ---
def generate_audios_for_scenes(scenes):
    # ... (Same as your previous generate_audios_for_scenes function)
    tts_model, prompt_tokenizer, desc_tokenizer = load_tts_model_and_tokenizers(TTS_MODEL_ID)
    audio_dir = os.path.join(get_session_temp_dir(), "audio_files")
    os.makedirs(audio_dir, exist_ok=True)
    audio_paths = []
    cols = st.columns(min(3, len(scenes)))
    col_idx = 0
    tts_description = "A neutral and clear narrator voice."
    for i, scene_text in enumerate(scenes):
        with st.spinner(f"Generating audio for scene {i+1}..."):
            try:
                desc_ids = desc_tokenizer(tts_description, return_tensors="pt").input_ids.to(tts_model.device)
                prompt_ids = prompt_tokenizer(scene_text, return_tensors="pt").input_ids.to(tts_model.device)
                generation_output = tts_model.generate(input_ids=desc_ids, prompt_input_ids=prompt_ids)
                audio_waveform = generation_output.to(torch.float32).cpu().numpy()
                file_path = os.path.join(audio_dir, f"audio_scene_{i+1}.wav")
                sf.write(file_path, audio_waveform, tts_model.config.sampling_rate)
                audio_paths.append(file_path)
                with cols[col_idx % len(cols)]:
                    st.markdown(f"**Audio Scene {i+1}**"); st.audio(file_path)
                    with open(file_path, "rb") as f_audio:
                        st.download_button(
                            f"Download Audio {i+1}", f_audio.read(), f"s_{i+1}_audio.wav", "audio/wav", key=f"dl_aud_{i}"
                        )
                col_idx +=1
            except Exception as e: st.error(f"Audio error s{i+1}: {e}"); audio_paths.append(None)
    clear_torch_cache()
    return audio_paths


# --- Step 4: Create Video (Keep from previous version) ---
def create_video_from_scenes(image_file_paths, audio_file_paths, output_filename="final_pov_video.mp4"):
    # ... (Same as your previous create_video_from_scenes function)
    if not image_file_paths or not audio_file_paths or len(image_file_paths) != len(audio_file_paths):
        st.error("Mismatch/missing assets for video."); return None
    try: subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
    except: st.error("FFMPEG not found."); return None

    temp_clips_dir = os.path.join(get_session_temp_dir(), "temp_video_clips")
    os.makedirs(temp_clips_dir, exist_ok=True)
    clips_paths, valid_scenes = [], 0
    for i, (img_p, aud_p) in enumerate(zip(image_file_paths, audio_file_paths)):
        if not (img_p and aud_p): continue
        try:
            aud_info = sf.info(aud_p); aud_dur = aud_info.duration
            if aud_dur <= 0.1: aud_dur = 1.0
            clip_p = os.path.join(temp_clips_dir, f"c_{i:03d}.mp4")
            cmd = [
                "ffmpeg", "-y", "-loop", "1", "-i", img_p, "-i", aud_p, "-c:v", "libx264", 
                "-preset", "medium", "-tune", "stillimage", "-c:a", "aac", "-b:a", "192k", 
                "-pix_fmt", "yuv420p", "-t", str(aud_dur), "-shortest", clip_p
            ]
            res = subprocess.run(cmd, capture_output=True, text=True)
            if res.returncode != 0: st.error(f"FFMPEG clip {i+1} err:\n{res.stderr}"); continue
            clips_paths.append(clip_p); valid_scenes +=1
        except Exception as e: st.error(f"Video scene {i+1} err: {e}")
    
    if not clips_paths or valid_scenes == 0: st.error("No valid video clips."); return None
    
    concat_list_f = os.path.join(temp_clips_dir, "concat_list.txt")
    with open(concat_list_f, "w") as f:
        for cp in clips_paths: f.write(f"file '{os.path.basename(cp)}'\n")
    final_vid_p = os.path.join(get_session_temp_dir(), output_filename)
    concat_cmd = ["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", concat_list_f, "-c", "copy", final_vid_p]
    res = subprocess.run(concat_cmd, capture_output=True, text=True, cwd=temp_clips_dir)
    if res.returncode != 0: st.error(f"FFMPEG concat err:\n{res.stderr}"); return None
    st.success("Video created!"); return final_vid_p

# --- Main App UI ---
st.title("⚡ POV Story Generator Pro (Gemini SDK Images)")
st.markdown("Create POV stories: Local LLM/TTS & Google Gemini SDK for Image Generation.")

if not google_gemini_sdk_available:
    st.error("Google Generative AI SDK (`google-generativeai`) is not installed. Image generation will not be available.")
    st.markdown("Please install it: `pip install google-generativeai`")
if not GEMINI_API_KEY:
     st.warning("`GEMINI_API_KEY` not set in secrets. Please configure it for Gemini SDK image generation.")

with st.sidebar:
    st.header("📝 Story Configuration")
    prompt = st.text_area(
        "POV story prompt:", st.session_state.get("user_prompt", "A squirrel discovers a tiny spaceship."),
        height=100, key="user_prompt_input"
    )
    num_scenes = st.slider("Number of Scenes:", 1, 8, st.session_state.get("num_scenes_val", 2), key="num_scenes_slider")

    st.header("🎨 Image Generation (Gemini SDK)")
    st.caption(f"Using Gemini Model: `{GEMINI_IMAGE_MODEL_ID}`")
    st.warning("Note: The specified Gemini model for images may be experimental. Image dimensions (width/height) are determined by the model.")
    
    # Width/Height sliders are kept for conceptual consistency but noted as not directly controlling this Gemini API.
    img_width_ui = st.slider("Image Width (Informational):", 512, 1536, st.session_state.get("img_width_val", 1024), 128, key="img_w_ui")
    img_height_ui = st.slider("Image Height (Informational):", 512, 1536, st.session_state.get("img_height_val", 1024), 128, key="img_h_ui")
    
    cols_seed = st.columns(2)
    with cols_seed[0]:
        base_seed_img = st.number_input("Base Image Seed (hint):", 0, 9999999, st.session_state.get("base_seed_val", 12345), key="base_seed_in")
    with cols_seed[1]:
        randomize_seed_img = st.checkbox("Randomize Seed (hint)", st.session_state.get("random_seed_bool", True), key="rand_seed_chk")

    st.markdown("---")
    can_generate = google_gemini_sdk_available and GEMINI_API_KEY
    if st.button("🚀 Generate Full Story & Assets", type="primary", use_container_width=True, disabled=not can_generate):
        st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        cleanup_temp_files()
        st.session_state.story_data = None; st.session_state.pil_images = None
        st.session_state.image_paths_for_video = None; st.session_state.audio_paths = None
        st.session_state.video_path = None
        st.session_state.user_prompt = prompt; st.session_state.num_scenes_val = num_scenes
        # Store informational width/height, though not used directly in Gemini SDK call here
        st.session_state.img_width_val = img_width_ui; st.session_state.img_height_val = img_height_ui 
        st.session_state.base_seed_val = base_seed_img; st.session_state.random_seed_bool = randomize_seed_img
        st.session_state.generate_all = True
    elif not can_generate:
        st.warning("Image generation disabled: missing Gemini API Key or SDK.")

    st.markdown("---")
    st.header("🛠️ Utilities")
    if st.button("🧹 Clear Cache & Temp Files & Restart", use_container_width=True):
        st.cache_resource.clear()
        keys_to_clear = ['story_data', 'pil_images', 'image_paths_for_video', 
                         'audio_paths', 'video_path', 'temp_base_dir', 'generate_all']
        for key in keys_to_clear:
            if key in st.session_state: del st.session_state[key]
        cleanup_temp_files(); st.session_state.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        st.success("Caches & temp files cleared. App restarting."); st.rerun()

# Main content area - (generation logic similar to previous, adapted for new image function)
if st.session_state.get("generate_all"):
    # --- 1. Story ---
    with st.status("🧠 Generating story...", True) as s_story:
        try:
            st.session_state.story_data = generate_story(st.session_state.user_prompt, st.session_state.num_scenes_val)
            if st.session_state.story_data: s_story.update(label="Story generated!", state="complete")
            else: s_story.update(label="Story fail.", state="error"); st.session_state.generate_all = False
        except Exception as e: st.error(f"Story err: {e}"); s_story.update(label="Story err.",state="error"); st.session_state.generate_all=False
    if st.session_state.story_data:
        st.subheader(f"🎬 Story: {st.session_state.story_data['title']}")
        for i, scene_text in enumerate(st.session_state.story_data['scenes']): st.markdown(f"**Scene {i+1}:** {scene_text}")
        story_json = json.dumps(st.session_state.story_data, indent=2)
        st.download_button("Download Story (JSON)", story_json, f"story.json", "application/json")
        st.markdown("---")

    # --- 2. Images ---
    if st.session_state.get("generate_all") and st.session_state.story_data:
        with st.status(f"🎨 Generating images via Gemini SDK ({GEMINI_IMAGE_MODEL_ID})...", True) as s_images:
            try:
                st.session_state.pil_images, st.session_state.image_paths_for_video = generate_images_for_scenes(
                    st.session_state.story_data['scenes'], 
                    st.session_state.random_seed_bool, st.session_state.base_seed_val
                )
                if st.session_state.pil_images and any(img for img in st.session_state.pil_images):
                    s_images.update(label="Images generated!", state="complete")
                else: s_images.update(label="Image gen failed/no images.", state="error"); st.session_state.generate_all = False
            except Exception as e: st.error(f"Image gen err: {e}"); s_images.update(label="Image error.",state="error"); st.session_state.generate_all=False
        st.markdown("---")

    # --- 3. Audio ---
    if st.session_state.get("generate_all") and st.session_state.story_data:
        with st.status("🔊 Generating audio...", True) as s_audio:
            try:
                st.session_state.audio_paths = generate_audios_for_scenes(st.session_state.story_data['scenes'])
                if st.session_state.audio_paths and any(p for p in st.session_state.audio_paths):
                    s_audio.update(label="Audio generated!", state="complete")
                else: s_audio.update(label="Audio gen failed.", state="error"); st.session_state.generate_all = False
            except Exception as e: st.error(f"Audio err: {e}"); s_audio.update(label="Audio error.",state="error"); st.session_state.generate_all=False
        st.markdown("---")

    # --- 4. Video ---
    if st.session_state.get("generate_all") and st.session_state.image_paths_for_video and st.session_state.audio_paths:
        valid_assets = sum(1 for im,au in zip(st.session_state.image_paths_for_video, st.session_state.audio_paths) if im and au)
        if valid_assets > 0:
            with st.status("📹 Creating video...", True) as s_video:
                try:
                    st.session_state.video_path = create_video_from_scenes(
                        st.session_state.image_paths_for_video, st.session_state.audio_paths
                    )
                    if st.session_state.video__path: s_video.update(label="Video created!", state="complete")
                    else: s_video.update(label="Video creation failed.", state="error")
                except Exception as e: st.error(f"Video err: {e}"); s_video.update(label="Video error.", state="error")
            if st.session_state.video_path:
                st.subheader("🎞️ Final Video"); st.video(st.session_state.video_path)
                with open(st.session_state.video_path, "rb") as fv:
                    st.download_button("Download Video", fv.read(), os.path.basename(st.session_state.video_path), "video/mp4")
            st.markdown("---")
        else: st.warning("Not enough assets for video.")
    if "generate_all" in st.session_state: del st.session_state.generate_all
elif not st.session_state.get("user_prompt"):
    st.info("Configure story in sidebar & click 'Generate' to begin!")