Update app.py

app.py

@@ -4,12 +4,18 @@ import torch
 from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
 from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
 from diffusers.utils.export_utils import export_to_video
+from diffusers import Flux2KleinPipeline
 import gradio as gr
 import tempfile
 import numpy as np
 from PIL import Image
 import random
 import gc
+import cv2
+import os
+import base64
+from io import BytesIO
+from typing import List
 
 from torchao.quantization import quantize_
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
@@ -18,6 +24,7 @@ from torchao.quantization import Int8WeightOnlyConfig
 import aoti
 
 
+# ─── Model & dimension constants ───────────────────────────────────────────────
 MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
 
 MAX_DIM = 832
@@ -31,17 +38,22 @@ FIXED_FPS = 16
 MIN_FRAMES_MODEL = 8
 MAX_FRAMES_MODEL = 80
 
-MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
-MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
+MIN_DURATION = round(MIN_FRAMES_MODEL / FIXED_FPS, 1)
+MAX_DURATION = round(MAX_FRAMES_MODEL / FIXED_FPS, 1)
 
+device = "cuda"
 
-pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID,
-    transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
+# ─── Load Wan 2.2 I2V pipeline ────────────────────────────────────────────────
+pipe = WanImageToVideoPipeline.from_pretrained(
+    MODEL_ID,
+    transformer=WanTransformer3DModel.from_pretrained(
+        'cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
         subfolder='transformer',
         torch_dtype=torch.bfloat16,
         device_map='cuda',
     ),
-    transformer_2=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
+    transformer_2=WanTransformer3DModel.from_pretrained(
+        'cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
         subfolder='transformer_2',
         torch_dtype=torch.bfloat16,
         device_map='cuda',
@@ -50,15 +62,15 @@ pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID,
 ).to('cuda')
 
 pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy", 
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
+    "Kijai/WanVideo_comfy",
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
     adapter_name="lightx2v"
 )
 kwargs_lora = {}
 kwargs_lora["load_into_transformer_2"] = True
 pipe.load_lora_weights(
-    "Kijai/WanVideo_comfy", 
-    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
+    "Kijai/WanVideo_comfy",
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors",
     adapter_name="lightx2v_2", **kwargs_lora
 )
 pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
@@ -79,10 +91,21 @@ spaces.aoti_load(
     repo_id='cbensimon/WanTransformer3DModel-sm120-cu130-raa',
 )
 
+# ─── Load FLUX.2 Klein 4B upscaler pipeline ────────────────────────────────────
+print("Loading FLUX.2 Klein 4B model...")
+klein_pipe = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-4B",
+    torch_dtype=torch.bfloat16,
+).to(device)
+print("FLUX.2 Klein 4B loaded successfully.")
+
 
+# ─── Prompts ───────────────────────────────────────────────────────────────────
 default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
 default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
 
+
+# ─── Image & frame helpers ─────────────────────────────────────────────────────
 def resize_image(image: Image.Image) -> Image.Image:
     """
     Resizes an image to fit within the model's constraints, preserving aspect ratio as much as possible.
@@ -94,12 +117,12 @@ def resize_image(image: Image.Image) -> Image.Image:
         return image.resize((SQUARE_DIM, SQUARE_DIM), Image.LANCZOS)
 
     aspect_ratio = width / height
-    
-    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM 
-    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM 
+
+    MAX_ASPECT_RATIO = MAX_DIM / MIN_DIM
+    MIN_ASPECT_RATIO = MIN_DIM / MAX_DIM
 
     image_to_resize = image
-    
+
     if aspect_ratio > MAX_ASPECT_RATIO:
         # Very wide image -> crop width to fit 832x480 aspect ratio
         target_w, target_h = MAX_DIM, MIN_DIM
@@ -125,7 +148,7 @@ def resize_image(image: Image.Image) -> Image.Image:
 
     final_w = max(MIN_DIM, min(MAX_DIM, final_w))
     final_h = max(MIN_DIM, min(MAX_DIM, final_h))
-    
+
     return image_to_resize.resize((final_w, final_h), Image.LANCZOS)
 
 
@@ -137,6 +160,70 @@ def get_num_frames(duration_seconds: float):
     ))
 
 
+def extract_frames_from_video(video_path: str, duration_seconds: float) -> List[Image.Image]:
+    """
+    Extract one frame per whole second from the generated video.
+    E.g. 3.5 s → frames at ~0.5 s, ~1.5 s, ~2.5 s (midpoints of each second bucket).
+    Returns PIL images.
+    """
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or FIXED_FPS
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    total_secs = total_frames / fps
+
+    n_buckets = max(1, int(duration_seconds))
+    frames_out = []
+    for i in range(n_buckets):
+        t = i + 0.5
+        t = min(t, total_secs - 0.05)
+        frame_idx = int(t * fps)
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+        ret, frame = cap.read()
+        if ret:
+            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames_out.append(Image.fromarray(rgb))
+    cap.release()
+    return frames_out
+
+
+def update_dimensions_for_klein(image: Image.Image):
+    """Calculate dimensions for Flux2Klein, snapped to multiples of 16."""
+    w, h = image.size
+    scale = min(1024 / w, 1024 / h)
+    nw = int(w * scale)
+    nh = int(h * scale)
+    return (nw // 16) * 16, (nh // 16) * 16
+
+
+# ─── Upscale a single frame with FLUX.2 Klein 4B ──────────────────────────────
+@spaces.GPU
+def upscale_single_frame(frame: Image.Image) -> Image.Image:
+    """
+    Upscale a single PIL image using FLUX.2 Klein 4B img2img pipeline.
+    Uses: seed=0, randomize_seed=True, guidance_scale=1.0, steps=4
+    """
+    target_w, target_h = update_dimensions_for_klein(frame)
+
+    # Resize input to the target upscale dimensions
+    resized_input = frame.resize((target_w, target_h), Image.LANCZOS)
+
+    current_seed = random.randint(0, MAX_SEED)
+
+    result = klein_pipe(
+        prompt="high quality, ultra detailed, sharp focus, 8k resolution",
+        image=resized_input,
+        height=target_h,
+        width=target_w,
+        guidance_scale=1.0,
+        num_inference_steps=4,
+        strength=0.4,
+        generator=torch.Generator(device=device).manual_seed(current_seed),
+    ).images[0]
+
+    return result
+
+
+# ─── Video generation ──────────────────────────────────────────────────────────
 def get_duration(
     input_image,
     prompt,
@@ -157,62 +244,26 @@ def get_duration(
     step_duration = BASE_STEP_DURATION * factor ** 1.5
     return 10 + int(steps) * step_duration
 
+
 @spaces.GPU(duration=get_duration)
 def generate_video(
     input_image,
     prompt,
-    steps = 4,
+    steps=4,
     negative_prompt=default_negative_prompt,
-    duration_seconds = MAX_DURATION,
-    guidance_scale = 1,
-    guidance_scale_2 = 1,    
-    seed = 42,
-    randomize_seed = False,
+    duration_seconds=MAX_DURATION,
+    guidance_scale=1,
+    guidance_scale_2=1,
+    seed=42,
+    randomize_seed=False,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
     Generate a video from an input image using the Wan 2.2 14B I2V model with Lightning LoRA.
-    
-    This function takes an input image and generates a video animation based on the provided
-    prompt and parameters. It uses an FP8 qunatized Wan 2.2 14B Image-to-Video model in with Lightning LoRA
-    for fast generation in 4-8 steps.
-    
-    Args:
-        input_image (PIL.Image): The input image to animate. Will be resized to target dimensions.
-        prompt (str): Text prompt describing the desired animation or motion.
-        steps (int, optional): Number of inference steps. More steps = higher quality but slower.
-            Defaults to 4. Range: 1-30.
-        negative_prompt (str, optional): Negative prompt to avoid unwanted elements. 
-            Defaults to default_negative_prompt (contains unwanted visual artifacts).
-        duration_seconds (float, optional): Duration of the generated video in seconds.
-            Defaults to 2. Clamped between MIN_FRAMES_MODEL/FIXED_FPS and MAX_FRAMES_MODEL/FIXED_FPS.
-        guidance_scale (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        guidance_scale_2 (float, optional): Controls adherence to the prompt. Higher values = more adherence.
-            Defaults to 1.0. Range: 0.0-20.0.
-        seed (int, optional): Random seed for reproducible results. Defaults to 42.
-            Range: 0 to MAX_SEED (2147483647).
-        randomize_seed (bool, optional): Whether to use a random seed instead of the provided seed.
-            Defaults to False.
-        progress (gr.Progress, optional): Gradio progress tracker. Defaults to gr.Progress(track_tqdm=True).
-    
-    Returns:
-        tuple: A tuple containing:
-            - video_path (str): Path to the generated video file (.mp4)
-            - current_seed (int): The seed used for generation (useful when randomize_seed=True)
-    
-    Raises:
-        gr.Error: If input_image is None (no image uploaded).
-    
-    Note:
-        - Frame count is calculated as duration_seconds * FIXED_FPS (24)
-        - Output dimensions are adjusted to be multiples of MOD_VALUE (32)
-        - The function uses GPU acceleration via the @spaces.GPU decorator
-        - Generation time varies based on steps and duration (see get_duration function)
     """
     if input_image is None:
         raise gr.Error("Please upload an input image.")
-    
+
     num_frames = get_num_frames(duration_seconds)
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     resized_image = resize_image(input_image)
@@ -237,36 +288,238 @@ def generate_video(
 
     return video_path, current_seed
 
-with gr.Blocks() as demo:
-    gr.Markdown("# Fast 4 steps Wan 2.2 I2V (14B) with Lightning LoRA")
-    gr.Markdown("run Wan 2.2 in just 4-8 steps, with [Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU⚡️")
+
+# ─── Full pipeline: generate → extract frames → upscale → return card data ───
+def generate_and_upscale(
+    input_image,
+    prompt,
+    steps,
+    negative_prompt,
+    duration_seconds,
+    guidance_scale,
+    guidance_scale_2,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    1. Generate video
+    2. Extract 1 frame per second
+    3. Upscale each frame with FLUX.2 Klein 4B
+    4. Return video, seed, and list of (original_frame, upscaled_frame) pairs
+    """
+    # Step 1: Generate video
+    video_path, current_seed = generate_video(
+        input_image, prompt, steps, negative_prompt,
+        duration_seconds, guidance_scale, guidance_scale_2,
+        seed, randomize_seed, progress
+    )
+
+    # Step 2: Extract frames (1 per second)
+    frames = extract_frames_from_video(video_path, duration_seconds)
+
+    if not frames:
+        return video_path, current_seed, []
+
+    # Step 3: Upscale each frame
+    upscaled_pairs = []
+    for i, frame in enumerate(frames):
+        progress(i / len(frames), desc=f"Upscaling frame {i+1}/{len(frames)}...")
+        upscaled = upscale_single_frame(frame)
+        upscaled_pairs.append((frame, upscaled))
+
+    return video_path, current_seed, upscaled_pairs
+
+
+def pil_to_temp_path(img: Image.Image, prefix: str = "frame") -> str:
+    """Save a PIL image to a temp file and return the path."""
+    with tempfile.NamedTemporaryFile(suffix=".png", prefix=prefix + "_", delete=False) as f:
+        img.save(f, format="PNG")
+        return f.name
+
+
+def run_pipeline(
+    input_image, prompt, steps, negative_prompt,
+    duration_seconds, guidance_scale, guidance_scale_2,
+    seed, randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Orchestrate the full pipeline and return outputs for Gradio.
+    Returns:
+        - video_path
+        - seed
+        - list of ImageSlider pairs for up to MAX_CARDS cards
+        - list of upscaled image paths for download buttons
+    """
+    MAX_CARDS = 10  # maximum frames we support in UI
+
+    video_path, current_seed, upscaled_pairs = generate_and_upscale(
+        input_image, prompt, steps, negative_prompt,
+        duration_seconds, guidance_scale, guidance_scale_2,
+        seed, randomize_seed, progress,
+    )
+
+    # Build outputs for each card slot
+    slider_outputs = []
+    download_outputs = []
+    visibility_outputs = []
+
+    for i in range(MAX_CARDS):
+        if i < len(upscaled_pairs):
+            orig, upscaled = upscaled_pairs[i]
+            orig_path = pil_to_temp_path(orig, prefix=f"original_sec{i+1}")
+            upscaled_path = pil_to_temp_path(upscaled, prefix=f"upscaled_sec{i+1}")
+            slider_outputs.append((orig_path, upscaled_path))
+            download_outputs.append(upscaled_path)
+            visibility_outputs.append(gr.update(visible=True))
+        else:
+            slider_outputs.append(None)
+            download_outputs.append(None)
+            visibility_outputs.append(gr.update(visible=False))
+
+    # Flatten: video, seed, then interleaved (slider, download, visibility) per card
+    results = [video_path, current_seed]
+    for i in range(MAX_CARDS):
+        results.append(slider_outputs[i])
+        results.append(download_outputs[i])
+        results.append(visibility_outputs[i])
+
+    return results
+
+
+# ─── Build Gradio UI ──────────────────────────────────────────────────────────
+
+css = '''
+.gradio-container .contain {
+    max-width: 1200px !important;
+    margin: 0 auto !important;
+}
+.upscale-card {
+    border: 1px solid var(--border-color-primary);
+    border-radius: 12px;
+    padding: 16px;
+    margin-bottom: 12px;
+    background: var(--background-fill-secondary);
+}
+.card-header {
+    font-size: 1.1em;
+    font-weight: 600;
+    margin-bottom: 8px;
+    color: var(--body-text-color);
+}
+'''
+
+MAX_CARDS = 10
+
+with gr.Blocks(css=css, title="wan2.2-i2v-fast") as demo:
+    gr.Markdown("# wan2.2-i2v-fast")
+    gr.Markdown(
+        "Run **Wan 2.2 I2V (14B)** in just 4-8 steps with "
+        "[Lightning LoRA](https://huggingface.co/Kijai/WanVideo_comfy/tree/main/Wan22-Lightning), "
+        "FP8 quantization & AoT compilation — compatible with 🧨 diffusers and ZeroGPU⚡️  \n"
+        "**+ FLUX.2 Klein 4B** frame upscaler — automatically extracts 1 frame/sec and upscales to ~1024px"
+    )
+
     with gr.Row():
-        with gr.Column():
+        # ─── LEFT COLUMN: Inputs ──────────────────────────────────
+        with gr.Column(scale=1):
             input_image_component = gr.Image(type="pil", label="Input Image")
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v, lines=3)
-            duration_seconds_input = gr.Slider(minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
-            
+            prompt_input = gr.Textbox(
+                label="Prompt", value=default_prompt_i2v, lines=3
+            )
+            duration_seconds_input = gr.Slider(
+                minimum=MIN_DURATION, maximum=MAX_DURATION, step=0.1, value=3.5,
+                label="Duration (seconds)",
+                info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps."
+            )
+
             with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
-                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
-                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                steps_slider = gr.Slider(minimum=1, maximum=30, step=1, value=6, label="Inference Steps") 
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale - high noise stage")
-                guidance_scale_2_input = gr.Slider(minimum=0.0, maximum=10.0, step=0.5, value=1, label="Guidance Scale 2 - low noise stage")
-
-            generate_button = gr.Button("Generate Video", variant="primary")
-        with gr.Column():
-            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-    
+                negative_prompt_input = gr.Textbox(
+                    label="Negative Prompt", value=default_negative_prompt, lines=3
+                )
+                seed_input = gr.Slider(
+                    label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42,
+                    interactive=True
+                )
+                randomize_seed_checkbox = gr.Checkbox(
+                    label="Randomize seed", value=True, interactive=True
+                )
+                steps_slider = gr.Slider(
+                    minimum=1, maximum=30, step=1, value=4,
+                    label="Inference Steps"
+                )
+                guidance_scale_input = gr.Slider(
+                    minimum=0.0, maximum=10.0, step=0.5, value=1,
+                    label="Guidance Scale - high noise stage"
+                )
+                guidance_scale_2_input = gr.Slider(
+                    minimum=0.0, maximum=10.0, step=0.5, value=1,
+                    label="Guidance Scale 2 - low noise stage"
+                )
+
+            generate_button = gr.Button("🎬 Generate Video & Upscale Frames", variant="primary")
+
+        # ─── RIGHT COLUMN: Video + Upscale Cards ──────────────────
+        with gr.Column(scale=1):
+            video_output = gr.Video(
+                label="Generated Video", autoplay=True, interactive=False
+            )
+
+            gr.Markdown("### 🔍 Upscaled Frame Comparisons")
+            gr.Markdown(
+                "*Drag the slider on each card to compare the **original frame** (left) "
+                "vs the **FLUX.2 Klein 4B upscaled** version (right). "
+                "Click download to save the upscaled image.*"
+            )
+
+            # Create card slots dynamically
+            card_groups = []   # list of (group, slider, download_btn)
+            slider_components = []
+            download_components = []
+            group_components = []
+
+            for i in range(MAX_CARDS):
+                with gr.Group(visible=False, elem_classes="upscale-card") as card_group:
+                    gr.Markdown(f"**Frame {i+1}** — Second {i+1}", elem_classes="card-header")
+                    img_slider = gr.ImageSlider(
+                        label=f"Original ↔ Upscaled (Second {i+1})",
+                        type="filepath",
+                        interactive=True,
+                    )
+                    download_btn = gr.File(
+                        label=f"⬇️ Download Upscaled Frame {i+1}",
+                        interactive=False,
+                    )
+
+                slider_components.append(img_slider)
+                download_components.append(download_btn)
+                group_components.append(card_group)
+
+    # ─── Wire up inputs/outputs ────────────────────────────────────────────────
     ui_inputs = [
         input_image_component, prompt_input, steps_slider,
         negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, guidance_scale_2_input, seed_input, randomize_seed_checkbox
+        guidance_scale_input, guidance_scale_2_input,
+        seed_input, randomize_seed_checkbox,
     ]
-    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
 
+    # Build output list: video, seed, then (slider, download, group_visibility) × MAX_CARDS
+    ui_outputs = [video_output, seed_input]
+    for i in range(MAX_CARDS):
+        ui_outputs.append(slider_components[i])
+        ui_outputs.append(download_components[i])
+        ui_outputs.append(group_components[i])
+
+    generate_button.click(
+        fn=run_pipeline,
+        inputs=ui_inputs,
+        outputs=ui_outputs,
+    )
+
+    # ─── Examples ──────────────────────────────────────────────────────────────
     gr.Examples(
-        examples=[ 
+        examples=[
             [
                 "wan_i2v_input.JPG",
                 "POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind (clear water, green hills, blue sky with clouds). Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles and sunlight beams. Brief underwater view of cat's face, then cat resurfaces, still filming selfie, playful summer vacation mood.",
@@ -283,12 +536,12 @@ with gr.Blocks() as demo:
                 6,
             ],
         ],
-        inputs=[input_image_component, prompt_input, steps_slider], outputs=[video_output, seed_input], fn=generate_video, cache_examples=True, cache_mode="lazy"
+        inputs=[input_image_component, prompt_input, steps_slider],
+        outputs=[video_output, seed_input],
+        fn=generate_video,
+        cache_examples=True,
+        cache_mode="lazy",
     )
 
-
-css = '''
-.gradio-container .contain{max-width: 1000px !important; margin: 0 auto !important}
-'''
 if __name__ == "__main__":
     demo.queue().launch(mcp_server=True, css=css)