Update app.py

app.py

@@ -2,7 +2,7 @@
 #
 # Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
 #
-# Version: 2.0.0
+# Version: 2.0.1
 #
 # Contact:
 # Carlos Rodrigues dos Santos
@@ -82,12 +82,12 @@ try:
         i18n = json.load(f)
 except Exception as e:
     logger.error(f"Error loading i18n.json: {e}")
-    i18n = {"🇧🇷": {}, "🇺🇸": {}, "🇨🇳": {}}
+    i18n = {"pt": {}, "en": {}, "zh": {}}
 
 # Fallback for missing languages
-if '🇧🇷' not in i18n: i18n['🇧🇷'] = i18n.get('🇺🇸', {})
-if '🇺🇸' not in i18n: i18n['🇺🇸'] = {}
-if '🇨🇳' not in i18n: i18n['🇨🇳'] = i18n.get('🇺🇸', {})
+if 'pt' not in i18n: i18n['pt'] = i18n.get('en', {})
+if 'en' not in i18n: i18n['en'] = {}
+if 'zh' not in i18n: i18n['zh'] = i18n.get('en', {})
 
 # Initialize the main orchestrator from the configuration file
 try:
@@ -105,10 +105,6 @@ except Exception as e:
 # the UI state after each operation.
 
 def run_pre_production_wrapper(prompt, num_keyframes, ref_files, resolution_str, duration_per_fragment, progress=gr.Progress()):
-    """
-    Wrapper for Pre-Production (Steps 1 & 2): Generates storyboard and keyframes.
-    This corresponds to the "Art Director Mode".
-    """
     if not ref_files:
         raise gr.Error("Please provide at least one reference image.")
 
@@ -119,10 +115,9 @@ def run_pre_production_wrapper(prompt, num_keyframes, ref_files, resolution_str,
 
     resolution = int(resolution_str.split('x')[0])
 
-    # Callback factory to create progress updates for keyframe generation
     def cb_factory(scene_index, total_scenes):
         start_time = time.time()
-        total_steps = 12 # Standard steps for Flux model
+        total_steps = 12
         def callback(pipe_self, step, timestep, callback_kwargs):
             elapsed = time.time() - start_time
             current_step = step + 1
@@ -138,14 +133,9 @@ def run_pre_production_wrapper(prompt, num_keyframes, ref_files, resolution_str,
 
     final_keyframes = aduc.task_generate_keyframes(storyboard, initial_ref_path, prompt, resolution, cb_factory)
 
-    # Make the next step (Production) visible
     return gr.update(value=storyboard), gr.update(value=final_keyframes), gr.update(visible=True, open=True)
 
 def run_pre_production_photo_wrapper(prompt, num_keyframes, ref_files, progress=gr.Progress()):
-    """
-    Wrapper for Pre-Production (Steps 1 & 2) in "Photographer Mode".
-    Generates a storyboard and selects the best matching keyframes from a user-provided pool.
-    """
     if not ref_files or len(ref_files) < 2:
         raise gr.Error("Photographer Mode requires at least 2 images: one base and one for the scene pool.")
 
@@ -165,19 +155,15 @@ def run_original_production_wrapper(keyframes, prompt, duration,
                                      guidance_scale, stg_scale, inference_steps,
                                      video_resolution,
                                      progress=gr.Progress()):
-    """
-    Wrapper for Step 3: Production. Generates the original master video using LTX.
-    Yields UI updates to show progress and final output.
-    """
+    """Wrapper for Step 3: Production. Correctly handles the return dictionary."""
     yield {
         original_video_output: gr.update(value=None, visible=True, label="🎬 Producing your original master video... Please wait."),
         final_video_output: gr.update(value=None, visible=True, label="🎬 Production in progress..."),
-        step4_accordion: gr.update(visible=False) # Hide post-production until this is done
+        step4_accordion: gr.update(visible=False)
     }
 
     resolution = int(video_resolution.split('x')[0])
     
-    # The orchestrator now returns the paths to the generated artifacts
     result = aduc.task_produce_original_movie(
         keyframes, prompt, duration,
         int(trim_percent), handler_strength, destination_convergence_strength,
@@ -191,17 +177,14 @@ def run_original_production_wrapper(keyframes, prompt, duration,
     yield {
         original_video_output: gr.update(value=original_video, label="✅ Original Master Video"),
         final_video_output: gr.update(value=original_video, label="Final Film (Result of the Last Step)"),
-        step4_accordion: gr.update(visible=True, open=True), # Show post-production tools
-        # Update state for the next steps
+        step4_accordion: gr.update(visible=True, open=True),
         original_latents_paths_state: original_latents,
         original_video_path_state: original_video,
         current_source_video_state: original_video,
     }
 
 def run_upscaler_wrapper(latent_paths, chunk_size, progress=gr.Progress()):
-    """
-    Wrapper for Post-Production Step 4A: Latent Upscaler.
-    """
+    """Wrapper for Step 4A: Latent Upscaler. Correctly handles the generator."""
     if not latent_paths:
         raise gr.Error("Cannot run Upscaler. No original latents found. Please complete Step 3 first.")
         
@@ -210,22 +193,19 @@ def run_upscaler_wrapper(latent_paths, chunk_size, progress=gr.Progress()):
         final_video_output: gr.update(label="Post-Production in progress: Latent Upscaling...")
     }
 
-    upscaled_video_path = aduc.task_run_latent_upscaler(
-        latent_paths, int(chunk_size), progress=progress
-    )
+    final_path = None
+    for path in aduc.task_run_latent_upscaler(latent_paths, int(chunk_size), progress=progress):
+        final_path = path['final_path']
 
     yield {
-        upscaler_video_output: gr.update(value=upscaled_video_path, label="✅ Latent Upscale Complete"),
-        final_video_output: gr.update(value=upscaled_video_path),
-        # Update states for subsequent steps
-        upscaled_video_path_state: upscaled_video_path,
-        current_source_video_state: upscaled_video_path,
+        upscaler_video_output: gr.update(value=final_path, label="✅ Latent Upscale Complete"),
+        final_video_output: gr.update(value=final_path),
+        upscaled_video_path_state: final_path,
+        current_source_video_state: final_path,
     }
 
-def run_hd_wrapper(source_video, model_version, steps, progress=gr.Progress()):
-    """
-    Wrapper for Post-Production Step 4B: HD Mastering.
-    """
+def run_hd_wrapper(source_video, model_version, steps, global_prompt, progress=gr.Progress()):
+    """Wrapper for Step 4B: HD Mastering. Correctly handles the generator."""
     if not source_video:
         raise gr.Error("Cannot run HD Mastering. No source video found. Please complete a previous step first.")
 
@@ -234,21 +214,19 @@ def run_hd_wrapper(source_video, model_version, steps, progress=gr.Progress()):
         final_video_output: gr.update(label="Post-Production in progress: HD Mastering...")
     }
     
-    hd_video_path = aduc.task_run_hd_mastering(
-        source_video, model_version, int(steps), progress=progress
-    )
+    final_path = None
+    for path in aduc.task_run_hd_mastering(source_video, model_version, int(steps), global_prompt, progress=progress):
+        final_path = path['final_path']
 
     yield {
-        hd_video_output: gr.update(value=hd_video_path, label="✅ HD Mastering Complete"),
-        final_video_output: gr.update(value=hd_video_path),
-        hd_video_path_state: hd_video_path,
-        current_source_video_state: hd_video_path,
+        hd_video_output: gr.update(value=final_path, label="✅ HD Mastering Complete"),
+        final_video_output: gr.update(value=final_path),
+        hd_video_path_state: final_path,
+        current_source_video_state: final_path,
     }
 
 def run_audio_wrapper(source_video, audio_prompt, global_prompt, progress=gr.Progress()):
-    """
-    Wrapper for Post-Production Step 4C: Audio Generation.
-    """
+    """Wrapper for Step 4C: Audio Generation. Correctly handles the generator."""
     if not source_video:
         raise gr.Error("Cannot run Audio Generation. No source video found. Please complete a previous step first.")
 
@@ -257,42 +235,32 @@ def run_audio_wrapper(source_video, audio_prompt, global_prompt, progress=gr.Pro
         final_video_output: gr.update(label="Post-Production in progress: Audio Generation...")
     }
 
-    # Use the specific audio prompt if provided, otherwise fall back to the global prompt
     final_audio_prompt = audio_prompt if audio_prompt and audio_prompt.strip() else global_prompt
 
-    video_with_audio_path = aduc.task_run_audio_generation(
-        source_video, final_audio_prompt, progress=progress
-    )
+    final_path = None
+    for path in aduc.task_run_audio_generation(source_video, final_audio_prompt, progress=progress):
+        final_path = path['final_path']
 
     yield {
-        audio_video_output: gr.update(value=video_with_audio_path, label="✅ Audio Generation Complete"),
-        final_video_output: gr.update(value=video_with_audio_path),
+        audio_video_output: gr.update(value=final_path, label="✅ Audio Generation Complete"),
+        final_video_output: gr.update(value=final_path),
     }
 
 def get_log_content():
-    """
-    Reads and returns the content of the log file for display in the UI.
-    """
     try:
         with open(LOG_FILE_PATH, "r", encoding="utf-8") as f:
             return f.read()
     except FileNotFoundError:
         return "Log file not yet created. Start a generation."
 
-def update_ui_language(lang_code):
-    """
-    Updates all text components in the UI to the selected language.
-    It fetches the translation map from the `i18n` dictionary.
-    """
+def update_ui_language(lang_emoji):
+    lang_code_map = {"🇧🇷": "pt", "🇺🇸": "en", "🇨🇳": "zh"}
+    lang_code = lang_code_map.get(lang_emoji, "en")
     lang_map = i18n.get(lang_code, i18n.get('en', {}))
-    # This dictionary maps each UI component variable to its new value from the language map.
     return {
-        # General
         title_md: gr.update(value=f"# {lang_map.get('app_title')}"),
         subtitle_md: gr.update(value=lang_map.get('app_subtitle')),
         lang_selector: gr.update(label=lang_map.get('lang_selector_label')),
-
-        # Step 1: Pre-Production
         step1_accordion: gr.update(label=lang_map.get('step1_accordion')),
         prompt_input: gr.update(label=lang_map.get('prompt_label')),
         ref_image_input: gr.update(label=lang_map.get('ref_images_label')),
@@ -303,11 +271,10 @@ def update_ui_language(lang_code):
         step1_mode_b_info_md: gr.update(value=f"*{lang_map.get('step1_mode_b_info')}*"),
         storyboard_output: gr.update(label=lang_map.get('storyboard_output_label')),
         keyframe_gallery: gr.update(label=lang_map.get('keyframes_gallery_label')),
-
-        # Step 3: Production
         step3_accordion: gr.update(label=lang_map.get('step3_accordion')),
         step3_description_md: gr.update(value=lang_map.get('step3_description')),
         produce_original_button: gr.update(value=lang_map.get('produce_original_button')),
+        ltx_advanced_options_accordion: gr.update(label=lang_map.get('ltx_advanced_options')),
         causality_accordion: gr.update(label=lang_map.get('causality_controls_title')),
         trim_percent_slider: gr.update(label=lang_map.get('trim_percent_label'), info=lang_map.get('trim_percent_info')),
         forca_guia_slider: gr.update(label=lang_map.get('forca_guia_label'), info=lang_map.get('forca_guia_info')),
@@ -316,8 +283,6 @@ def update_ui_language(lang_code):
         guidance_scale_slider: gr.update(label=lang_map.get('guidance_scale_label'), info=lang_map.get('guidance_scale_info')),
         stg_scale_slider: gr.update(label=lang_map.get('stg_scale_label'), info=lang_map.get('stg_scale_info')),
         inference_steps_slider: gr.update(label=lang_map.get('steps_label'), info=lang_map.get('steps_info')),
-
-        # Step 4: Post-Production
         step4_accordion: gr.update(label=lang_map.get('step4_accordion')),
         step4_description_md: gr.update(value=lang_map.get('step4_description')),
         sub_step_a_accordion: gr.update(label=lang_map.get('sub_step_a_upscaler')),
@@ -336,8 +301,6 @@ def update_ui_language(lang_code):
         audio_options_accordion: gr.update(label=lang_map.get('audio_options')),
         audio_prompt_input: gr.update(label=lang_map.get('audio_prompt_label'), info=lang_map.get('audio_prompt_info')),
         run_audio_button: gr.update(value=lang_map.get('run_audio_button')),
-
-        # Final Outputs & Logs
         final_video_output: gr.update(label=lang_map.get('final_video_label')),
         log_accordion: gr.update(label=lang_map.get('log_accordion_label')),
         log_display: gr.update(label=lang_map.get('log_display_label')),
@@ -346,25 +309,23 @@ def update_ui_language(lang_code):
 
 # --- 3. GRADIO UI DEFINITION ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # Initialize UI with default language (Portuguese)
     default_lang = i18n.get('pt', {})
-
-    # State components to manage the pipeline artifacts
+    
+    # State components
     original_latents_paths_state = gr.State(value=None)
     original_video_path_state = gr.State(value=None)
     upscaled_video_path_state = gr.State(value=None)
     hd_video_path_state = gr.State(value=None)
-    current_source_video_state = gr.State(value=None) # Tracks the latest video for post-production steps
+    current_source_video_state = gr.State(value=None)
 
-    # --- UI Header ---
+    # UI Header
     title_md = gr.Markdown(f"# {default_lang.get('app_title')}")
     subtitle_md = gr.Markdown(default_lang.get('app_subtitle'))
-
     with gr.Row():
-        lang_selector = gr.Radio(["🇧🇷", "🇺🇸", "🇨🇳"], value="pt", label=default_lang.get('lang_selector_label'))
+        lang_selector = gr.Radio(["🇧🇷", "🇺🇸", "🇨🇳"], value="🇧🇷", label=default_lang.get('lang_selector_label'))
         resolution_selector = gr.Radio(["480x480", "720x720", "960x960"], value="480x480", label="Base Resolution")
 
-    # --- Step 1 & 2: Pre-Production ---
+    # Step 1 & 2: Pre-Production
     with gr.Accordion(default_lang.get('step1_accordion'), open=True) as step1_accordion:
         prompt_input = gr.Textbox(label=default_lang.get('prompt_label'), value="A majestic lion walks across the savanna, sits down, and then roars at the setting sun.")
         ref_image_input = gr.File(label=default_lang.get('ref_images_label'), file_count="multiple", file_types=["image"])
@@ -378,112 +339,75 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         storyboard_output = gr.JSON(label=default_lang.get('storyboard_output_label'))
         keyframe_gallery = gr.Gallery(label=default_lang.get('keyframes_gallery_label'), visible=True, object_fit="contain", height="auto", type="filepath")
 
-    # --- Step 3: Production ---
+    # Step 3: Production
     with gr.Accordion(default_lang.get('step3_accordion'), open=False, visible=False) as step3_accordion:
         step3_description_md = gr.Markdown(default_lang.get('step3_description'))
-
         with gr.Accordion(default_lang.get('ltx_advanced_options'), open=False) as ltx_advanced_options_accordion:
             with gr.Accordion(default_lang.get('causality_controls_title'), open=True) as causality_accordion:
                 trim_percent_slider = gr.Slider(minimum=10, maximum=90, value=50, step=5, label=default_lang.get('trim_percent_label'), info=default_lang.get('trim_percent_info'))
                 with gr.Row():
                     forca_guia_slider = gr.Slider(label=default_lang.get('forca_guia_label'), minimum=0.0, maximum=1.0, value=0.5, step=0.05, info=default_lang.get('forca_guia_info'))
                     convergencia_destino_slider = gr.Slider(label=default_lang.get('convergencia_final_label'), minimum=0.0, maximum=1.0, value=0.75, step=0.05, info=default_lang.get('convergencia_final_info'))
-
             with gr.Accordion(default_lang.get('ltx_pipeline_options'), open=True) as ltx_pipeline_accordion:
                 with gr.Row():
                     guidance_scale_slider = gr.Slider(minimum=1.0, maximum=10.0, value=2.0, step=0.1, label=default_lang.get('guidance_scale_label'), info=default_lang.get('guidance_scale_info'))
                     stg_scale_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.025, step=0.005, label=default_lang.get('stg_scale_label'), info=default_lang.get('stg_scale_info'))
                 inference_steps_slider = gr.Slider(minimum=10, maximum=50, value=20, step=1, label=default_lang.get('steps_label'), info=default_lang.get('steps_info'))
-
         produce_original_button = gr.Button(default_lang.get('produce_original_button'), variant="primary")
-        original_video_output = gr.Video(label="Original Master Video", visible=False)
+        original_video_output = gr.Video(label="Original Master Video", visible=False, interactive=False)
 
-    # --- Step 4: Post-Production ---
+    # Step 4: Post-Production
     with gr.Accordion(default_lang.get('step4_accordion'), open=False, visible=False) as step4_accordion:
         step4_description_md = gr.Markdown(default_lang.get('step4_description'))
-
-        # Sub-Step 4A: Latent Upscaler
         with gr.Accordion(default_lang.get('sub_step_a_upscaler'), open=True) as sub_step_a_accordion:
             upscaler_description_md = gr.Markdown(default_lang.get('upscaler_description'))
             with gr.Accordion(default_lang.get('upscaler_options'), open=False) as upscaler_options_accordion:
-                upscaler_chunk_size_slider = gr.Slider(minimum=1, maximum=10, value=4, step=1, label=default_lang.get('upscaler_chunk_size_label'), info=default_lang.get('upscaler_chunk_size_info'))
+                upscaler_chunk_size_slider = gr.Slider(minimum=1, maximum=10, value=2, step=1, label=default_lang.get('upscaler_chunk_size_label'), info=default_lang.get('upscaler_chunk_size_info'))
             run_upscaler_button = gr.Button(default_lang.get('run_upscaler_button'), variant="secondary")
-            upscaler_video_output = gr.Video(label="Upscaled Video", visible=False)
-
-        # Sub-Step 4B: HD Mastering
+            upscaler_video_output = gr.Video(label="Upscaled Video", visible=False, interactive=False)
         with gr.Accordion(default_lang.get('sub_step_b_hd'), open=True) as sub_step_b_accordion:
             hd_description_md = gr.Markdown(default_lang.get('hd_description'))
             with gr.Accordion(default_lang.get('hd_options'), open=False) as hd_options_accordion:
-                hd_model_radio = gr.Radio(["3B", "7B"], value="3B", label=default_lang.get('hd_model_label'))
-                hd_steps_slider = gr.Slider(minimum=20, maximum=150, value=50, step=5, label=default_lang.get('hd_steps_label'), info=default_lang.get('hd_steps_info'))
+                hd_model_radio = gr.Radio(["3B", "7B"], value="7B", label=default_lang.get('hd_model_label'))
+                hd_steps_slider = gr.Slider(minimum=20, maximum=150, value=100, step=5, label=default_lang.get('hd_steps_label'), info=default_lang.get('hd_steps_info'))
             run_hd_button = gr.Button(default_lang.get('run_hd_button'), variant="secondary")
-            hd_video_output = gr.Video(label="HD Mastered Video", visible=False)
-
-        # Sub-Step 4C: Audio Generation
+            hd_video_output = gr.Video(label="HD Mastered Video", visible=False, interactive=False)
         with gr.Accordion(default_lang.get('sub_step_c_audio'), open=True) as sub_step_c_accordion:
             audio_description_md = gr.Markdown(default_lang.get('audio_description'))
             with gr.Accordion(default_lang.get('audio_options'), open=False) as audio_options_accordion:
                 audio_prompt_input = gr.Textbox(label=default_lang.get('audio_prompt_label'), info=default_lang.get('audio_prompt_info'), lines=3)
             run_audio_button = gr.Button(default_lang.get('run_audio_button'), variant="secondary")
-            audio_video_output = gr.Video(label="Video with Audio", visible=False)
-
-    # --- Final Output & Logs ---
-    final_video_output = gr.Video(label=default_lang.get('final_video_label'), visible=False)
+            audio_video_output = gr.Video(label="Video with Audio", visible=False, interactive=False)
 
+    # Final Output & Logs
+    final_video_output = gr.Video(label=default_lang.get('final_video_label'), visible=False, interactive=False)
     with gr.Accordion(default_lang.get('log_accordion_label'), open=False) as log_accordion:
         log_display = gr.Textbox(label=default_lang.get('log_display_label'), lines=20, interactive=False, autoscroll=True)
         update_log_button = gr.Button(default_lang.get('update_log_button'))
 
     # --- 4. UI EVENT CONNECTIONS ---
-    # Collect all UI components that need language updates
-    all_ui_components = list(update_ui_language('pt').keys())
+    all_ui_components = list(update_ui_language('🇧🇷').keys())
     lang_selector.change(fn=update_ui_language, inputs=lang_selector, outputs=all_ui_components)
 
-    # Pre-Production Button Clicks
-    storyboard_and_keyframes_button.click(
-        fn=run_pre_production_wrapper,
-        inputs=[prompt_input, num_keyframes_slider, ref_image_input, resolution_selector, duration_per_fragment_slider],
-        outputs=[storyboard_output, keyframe_gallery, step3_accordion]
-    )
-
-    storyboard_from_photos_button.click(
-        fn=run_pre_production_photo_wrapper,
-        inputs=[prompt_input, num_keyframes_slider, ref_image_input],
-        outputs=[storyboard_output, keyframe_gallery, step3_accordion]
-    )
+    storyboard_and_keyframes_button.click(fn=run_pre_production_wrapper, inputs=[prompt_input, num_keyframes_slider, ref_image_input, resolution_selector, duration_per_fragment_slider], outputs=[storyboard_output, keyframe_gallery, step3_accordion])
+    storyboard_from_photos_button.click(fn=run_pre_production_photo_wrapper, inputs=[prompt_input, num_keyframes_slider, ref_image_input], outputs=[storyboard_output, keyframe_gallery, step3_accordion])
 
-    # Production Button Click
     produce_original_button.click(
         fn=run_original_production_wrapper,
-        inputs=[
-            keyframe_gallery, prompt_input, duration_per_fragment_slider,
-            trim_percent_slider, forca_guia_slider, convergencia_destino_slider,
-            guidance_scale_slider, stg_scale_slider, inference_steps_slider,
-            resolution_selector
-        ],
-        outputs=[
-            original_video_output, final_video_output, step4_accordion,
-            original_latents_paths_state, original_video_path_state, current_source_video_state
-        ]
+        inputs=[keyframe_gallery, prompt_input, duration_per_fragment_slider, trim_percent_slider, forca_guia_slider, convergencia_destino_slider, guidance_scale_slider, stg_scale_slider, inference_steps_slider, resolution_selector],
+        outputs=[original_video_output, final_video_output, step4_accordion, original_latents_paths_state, original_video_path_state, current_source_video_state]
     )
     
-    # Post-Production Button Clicks
     run_upscaler_button.click(
         fn=run_upscaler_wrapper,
         inputs=[original_latents_paths_state, upscaler_chunk_size_slider],
-        outputs=[
-            upscaler_video_output, final_video_output,
-            upscaled_video_path_state, current_source_video_state
-        ]
+        outputs=[upscaler_video_output, final_video_output, upscaled_video_path_state, current_source_video_state]
     )
     
     run_hd_button.click(
         fn=run_hd_wrapper,
-        inputs=[current_source_video_state, hd_model_radio, hd_steps_slider],
-        outputs=[
-            hd_video_output, final_video_output,
-            hd_video_path_state, current_source_video_state
-        ]
+        inputs=[current_source_video_state, hd_model_radio, hd_steps_slider, prompt_input],
+        outputs=[hd_video_output, final_video_output, hd_video_path_state, current_source_video_state]
     )
 
     run_audio_button.click(
@@ -492,7 +416,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[audio_video_output, final_video_output]
     )
     
-    # Log Button Click
     update_log_button.click(fn=get_log_content, inputs=[], outputs=[log_display])
 
 # --- 5. APPLICATION LAUNCH ---