Update app.py

app.py

@@ -15,6 +15,7 @@ import torch
 # --- NEW ---
 # Import the OpenCV library
 import cv2
+import gc
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -37,6 +38,27 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import shutil
 
+MAX_SEED = np.iinfo(np.int32).max
+
+#import diffusers
+from diffusers import StableDiffusionXLImg2ImgPipeline, AutoencoderKL
+print("Loading SDXL Image-to-Image pipeline...")
+#vaeX = AutoencoderKL.from_pretrained('stabilityai/stable-diffusion-xl-refiner-1.0',subfolder='vae')
+enhancer_pipeline = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    #"stabilityai/stable-diffusion-xl-base-1.0",
+    "ford442/stable-diffusion-xl-refiner-1.0-bf16",
+    #torch_dtype=torch.bfloat16,
+    #variant="fp16",
+    #use_safetensors=True,
+    requires_aesthetics_score=True,
+    #vae=None
+)
+#enhancer_pipeline.vae=vaeX
+enhancer_pipeline.vae.set_default_attn_processor()
+
+enhancer_pipeline.to("cpu")
+print("SDXL Image-to-Image pipeline loaded successfully.")
+
 from inference import (
     create_ltx_video_pipeline,
     create_latent_upsampler,
@@ -87,58 +109,64 @@ def calculate_new_dimensions(orig_w, orig_h):
         new_w, new_h = 768, round((768 * (orig_h / orig_w)) / 32) * 32
     return int(max(256, min(new_h, MAX_IMAGE_SIZE))), int(max(256, min(new_w, MAX_IMAGE_SIZE)))
 
+
 def get_duration(*args, **kwargs):
     duration_ui = kwargs.get('duration_ui', 5.0)
-    if duration_ui > 20.0: return 120
-    if duration_ui > 13.0: return 90
-    if duration_ui > 7.0: return 75
-    if duration_ui > 5.0: return 60
-    return 45
+    if duration_ui > 7.0: return 120
+    if duration_ui > 5.0: return 100
+    return 90
+
+
+@spaces.GPU()
+def enhance_frame(image_to_enhance: Image.Image):
+    try:
+        print("Moving enhancer pipeline to GPU...")
+        seed = random.randint(0, MAX_SEED)
+        generator = torch.Generator(device='cpu').manual_seed(seed)
+        enhancer_pipeline.to("cuda",torch.bfloat16)
+        refine_prompt = "cinematic, high detail, sharp focus, 8k, professional photography"
+        enhanced_image = enhancer_pipeline(prompt=refine_prompt, image=image_to_enhance, strength=0.125, generator=generator, num_inference_steps=100).images[0]
+        print("Frame enhancement successful.")
+    except Exception as e:
+        print(f"Error during frame enhancement: {e}")
+        gr.Warning("Frame enhancement failed. Using original frame.")
+        return image_to_enhance
+    finally:
+        print("Moving enhancer pipeline to CPU...")
+        enhancer_pipeline.to("cpu")
+        gc.collect()
+        torch.cuda.empty_cache()
+    return enhanced_image
 
-# --- MODIFIED ---
-# This function is now completely rewritten to use the more robust OpenCV library
-# instead of imageio. This should definitively fix the FFmpeg error.
-def use_last_frame_as_input(video_filepath):
+
+def use_last_frame_as_input(video_filepath, do_enhance):
     if not video_filepath or not os.path.exists(video_filepath):
-        gr.Warning("No video available to get the last frame from.")
+        gr.Warning("No video clip available.")
         return None, gr.update()
-    
+    cap = None
     try:
-        print(f"Reading last frame from {video_filepath} using OpenCV...")
         cap = cv2.VideoCapture(video_filepath)
-        
-        # Get the total number of frames
         frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        if frame_count < 1:
-            raise ValueError("Video file could not be read or contains no frames.")
-            
-        # Set the position to the last frame (frame indices are 0-based)
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_count - 1)
-        
-        # Read the frame
         ret, frame = cap.read()
-        
-        if not ret or frame is None:
-            raise ValueError("Failed to read the last frame from the video.")
-            
-        # OpenCV reads frames in BGR format, so we convert it to RGB for saving with PIL
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        
-        pil_image = Image.fromarray(frame_rgb)
-        output_image_path = os.path.join(tempfile.mkdtemp(), f"last_frame_{random.randint(10000,99999)}.png")
-        pil_image.save(output_image_path)
-        
-        print("Successfully extracted last frame.")
-        return output_image_path, gr.update(selected="i2v_tab")
-
+        if not ret: raise ValueError("Failed to read frame.")
+        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                # 1. Immediately yield the original frame to the UI
+        print("Displaying original last frame...")
+        yield pil_image, gr.update()
+        if do_enhance:
+            enhanced_image = enhance_frame(pil_image)
+            # 2. Yield the enhanced frame and switch the tab
+            print("Displaying enhanced frame and switching tab...")
+            yield enhanced_image, gr.update(selected="i2v_tab")
+        else:
+            # If not enhancing, just switch the tab
+            yield pil_image, gr.update(selected="i2v_tab")
     except Exception as e:
-        print(f"Error extracting last frame with OpenCV: {e}")
-        gr.Error(f"Failed to extract the last frame: {e}")
+        gr.Error(f"Failed to extract frame: {e}")
         return None, gr.update()
     finally:
-        # Release the video capture object
-        if 'cap' in locals() and cap.isOpened():
-            cap.release()
+        if cap: cap.release()
 
 
 def stitch_videos(clips_list):
@@ -156,19 +184,22 @@ def stitch_videos(clips_list):
     except Exception as e:
         raise gr.Error(f"Failed to stitch videos: {e}")
 
+
 def clear_clips():
     return [], "Clips created: 0", None, None
 
+
 @spaces.GPU(duration=get_duration)
 def generate(prompt, negative_prompt, clips_list, input_image_filepath, input_video_filepath,
              height_ui, width_ui, mode, duration_ui, ui_frames_to_use,
              seed_ui, randomize_seed, ui_guidance_scale, improve_texture_flag, num_steps, fps,
              progress=gr.Progress(track_tqdm=True)):
-
-    # ... (the rest of the generate function is unchanged) ...
-    if mode not in ["text-to-video", "image-to-video", "video-to-video"]: raise gr.Error(f"Invalid mode: {mode}.")
-    if mode == "image-to-video" and not input_image_filepath: raise gr.Error("input_image_filepath is required for image-to-video mode")
-    elif mode == "video-to-video" and not input_video_filepath: raise gr.Error("input_video_filepath is required for video-to-video mode")
+    if mode not in ["text-to-video", "image-to-video", "video-to-video"]:
+        raise gr.Error(f"Invalid mode: {mode}.")
+    if mode == "image-to-video" and not input_image_filepath:
+        raise gr.Error("input_image_filepath is required for image-to-video mode")
+    elif mode == "video-to-video" and not input_video_filepath:
+        raise gr.Error("input_video_filepath is required for video-to-video mode")
     if randomize_seed: seed_ui = random.randint(0, 2**32 - 1)
     seed_everething(int(seed_ui))
     actual_num_frames = max(9, min(MAX_NUM_FRAMES, int(round((max(1, round(duration_ui * fps)) - 1.0) / 8.0) * 8 + 1)))
@@ -205,11 +236,16 @@ def generate(prompt, negative_prompt, clips_list, input_image_filepath, input_vi
     counter_text = f"Clips created: {len(updated_clips_list)}"
     return output_video_path, seed_ui, gr.update(visible=True), updated_clips_list, counter_text
 
-# ... (UI layout and other event handlers are unchanged) ...
-def update_task_image(): return "image-to-video"
-def update_task_text(): return "text-to-video"
-def update_task_video(): return "video-to-video"
+
+def update_task_image(): 
+    return "image-to-video"
+def update_task_text(): 
+    return "text-to-video"
+def update_task_video(): 
+    return "video-to-video"
+    
 css="""#col-container{margin:0 auto;max-width:900px;}"""
+
 with gr.Blocks(css=css) as demo:
     clips_state = gr.State([])
     gr.Markdown("# LTX Video Clip Stitcher")
@@ -218,13 +254,23 @@ with gr.Blocks(css=css) as demo:
         with gr.Column():
             with gr.Tabs() as tabs:
                 with gr.Tab("image-to-video", id="i2v_tab") as image_tab:
-                    video_i_hidden = gr.Textbox(visible=False); image_i2v = gr.Image(label="Input Image", type="filepath", sources=["upload", "webcam", "clipboard"]); i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3); i2v_button = gr.Button("Generate Image-to-Video Clip", variant="primary")
+                    video_i_hidden = gr.Textbox(visible=False); 
+                    image_i2v = gr.Image(label="Input Image", type="filepath", sources=["upload", "webcam", "clipboard"]); 
+                    i2v_prompt = gr.Textbox(label="Prompt", value="The creature from the image starts to move", lines=3);
+                    i2v_button = gr.Button("Generate Image-to-Video Clip", variant="primary")
                 with gr.Tab("text-to-video", id="t2v_tab") as text_tab:
-                    image_n_hidden = gr.Textbox(visible=False); video_n_hidden = gr.Textbox(visible=False); t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3); t2v_button = gr.Button("Generate Text-to-Video Clip", variant="primary")
+                    image_n_hidden = gr.Textbox(visible=False); 
+                    video_n_hidden = gr.Textbox(visible=False); t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3); 
+                    t2v_button = gr.Button("Generate Text-to-Video Clip", variant="primary")
                 with gr.Tab("video-to-video", id="v2v_tab") as video_tab:
-                    image_v_hidden = gr.Textbox(visible=False); video_v2v = gr.Video(label="Input Video", sources=["upload", "webcam"]); frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=120, value=9, step=8, info="Must be N*8+1."); v2v_prompt = gr.Textbox(label="Prompt", value="Change the style to cinematic anime", lines=3); v2v_button = gr.Button("Generate Video-to-Video Clip", variant="primary")
+                    image_v_hidden = gr.Textbox(visible=False); 
+                    video_v2v = gr.Video(label="Input Video", sources=["upload", "webcam"]); 
+                    frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=120, value=9, step=8, info="Must be N*8+1."); 
+                    v2v_prompt = gr.Textbox(label="Prompt", value="Change the style to cinematic anime", lines=3); 
+                    v2v_button = gr.Button("Generate Video-to-Video Clip", variant="primary")
             duration_input = gr.Slider(label="Clip Duration (seconds)", minimum=1.0, maximum=10.0, value=2.0, step=0.1)
             improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True)
+            enhance_checkbox = gr.Checkbox(label="Improve Frame (SDXL Refiner)", value=True)
         with gr.Column():
             output_video = gr.Video(label="Last Generated Clip", interactive=False)
             use_last_frame_button = gr.Button("Use Last Frame as Input Image", visible=False)
@@ -233,26 +279,40 @@ with gr.Blocks(css=css) as demo:
                 with gr.Row(): stitch_button = gr.Button("🎬 Stitch All Clips"); clear_button = gr.Button("🗑️ Clear All Clips")
                 final_video_output = gr.Video(label="Final Stitched Video", interactive=False)
     with gr.Accordion("Advanced settings", open=False):
-        mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False); negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", lines=2)
-        with gr.Row(): seed_input = gr.Number(label="Seed", value=42, precision=0); randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-        with gr.Row(visible=False): guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0), step=0.1)
-        with gr.Row(): height_input = gr.Slider(label="Height", value=768, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); width_input = gr.Slider(label="Width", value=768, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); num_steps = gr.Slider(label="Steps", value=20, step=1, minimum=1, maximum=420); fps = gr.Slider(label="FPS", value=30.0, step=1.0, minimum=4.0, maximum=60.0)
+        mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False); 
+        negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion, blurry, jittery, distorted", lines=2)
+        with gr.Row(): 
+            seed_input = gr.Number(label="Seed", value=42, precision=0); 
+            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
+        with gr.Row(visible=False): 
+            guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=PIPELINE_CONFIG_YAML.get("first_pass", {}).get("guidance_scale", 1.0), step=0.1)
+        with gr.Row(): 
+            height_input = gr.Slider(label="Height", value=768, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); 
+            width_input = gr.Slider(label="Width", value=768, step=32, minimum=32, maximum=MAX_IMAGE_SIZE); 
+        num_steps = gr.Slider(label="Steps", value=20, step=1, minimum=1, maximum=420); 
+        fps = gr.Slider(label="FPS", value=30.0, step=1.0, minimum=4.0, maximum=60.0)
     def handle_image_upload_for_dims(f, h, w):
         if not f: return gr.update(value=h), gr.update(value=w)
         img = Image.open(f); new_h, new_w = calculate_new_dimensions(img.width, img.height); return gr.update(value=new_h), gr.update(value=new_w)
     def handle_video_upload_for_dims(f, h, w):
         if not f or not os.path.exists(str(f)): return gr.update(value=h), gr.update(value=w)
         with imageio.get_reader(str(f)) as reader:
-            meta = reader.get_meta_data(); orig_w, orig_h = meta.get('size', (reader.get_data(0).shape[1], reader.get_data(0).shape[0])); new_h, new_w = calculate_new_dimensions(orig_w, orig_h); return gr.update(value=new_h), gr.update(value=new_w)
-    image_i2v.upload(handle_image_upload_for_dims, [image_i2v, height_input, width_input], [height_input, width_input]); video_v2v.upload(handle_video_upload_for_dims, [video_v2v, height_input, width_input], [height_input, width_input]); image_tab.select(update_task_image, outputs=[mode]); text_tab.select(update_task_text, outputs=[mode]); video_tab.select(update_task_video, outputs=[mode])
+            meta = reader.get_meta_data(); orig_w, orig_h = meta.get('size', (reader.get_data(0).shape[1], reader.get_data(0).shape[0])); 
+            new_h, new_w = calculate_new_dimensions(orig_w, orig_h); return gr.update(value=new_h), gr.update(value=new_w)
+    image_i2v.upload(handle_image_upload_for_dims, [image_i2v, height_input, width_input], [height_input, width_input]); 
+    video_v2v.upload(handle_video_upload_for_dims, [video_v2v, height_input, width_input], [height_input, width_input]); 
+    image_tab.select(update_task_image, outputs=[mode]); text_tab.select(update_task_text, outputs=[mode]); 
+    video_tab.select(update_task_video, outputs=[mode])
     common_params = [height_input, width_input, mode, duration_input, frames_to_use, seed_input, randomize_seed_input, guidance_scale_input, improve_texture, num_steps, fps]
-    t2v_inputs = [t2v_prompt, negative_prompt_input, clips_state, image_n_hidden, video_n_hidden] + common_params; i2v_inputs = [i2v_prompt, negative_prompt_input, clips_state, image_i2v, video_i_hidden] + common_params; v2v_inputs = [v2v_prompt, negative_prompt_input, clips_state, image_v_hidden, video_v2v] + common_params
+    t2v_inputs = [t2v_prompt, negative_prompt_input, clips_state, image_n_hidden, video_n_hidden] + common_params; 
+    i2v_inputs = [i2v_prompt, negative_prompt_input, clips_state, image_i2v, video_i_hidden] + common_params; 
+    v2v_inputs = [v2v_prompt, negative_prompt_input, clips_state, image_v_hidden, video_v2v] + common_params
     gen_outputs = [output_video, seed_input, use_last_frame_button, clips_state, clip_counter_display]
     hide_btn = lambda: gr.update(visible=False)
     t2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=t2v_inputs, outputs=gen_outputs, api_name="text_to_video")
     i2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=i2v_inputs, outputs=gen_outputs, api_name="image_to_video")
     v2v_button.click(hide_btn, outputs=[use_last_frame_button], queue=False).then(fn=generate, inputs=v2v_inputs, outputs=gen_outputs, api_name="video_to_video")
-    use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[output_video], outputs=[image_i2v, tabs])
+    use_last_frame_button.click(fn=use_last_frame_as_input, inputs=[output_video,enhance_checkbox], outputs=[image_i2v, tabs])
     stitch_button.click(fn=stitch_videos, inputs=[clips_state], outputs=[final_video_output])
     clear_button.click(fn=clear_clips, outputs=[clips_state, clip_counter_display, output_video, final_video_output])
 if __name__ == "__main__":