Spaces:

nockhtu
/

Wan-Extender

Paused

Stream and preserve intermediate video checkpoints

f2cc5c0 verified about 2 months ago

17.4 kB

from __future__ import annotations

import gc
import os
import random
import tempfile
import traceback
import zipfile

# ZeroGPU currently needs this compatible CUDA/PyTorch build before importing torch.
os.system(
    'pip install --quiet --upgrade --pre --extra-index-url '
    'https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces'
)

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import AutoPipelineForText2Image, FlowMatchEulerDiscreteScheduler
from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video
from PIL import Image
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig,
    quantize_,
)
from transformers import CLIPImageProcessor


WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
WAN_WEIGHTS_ID = "cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers"
IMAGE_MODEL_ID = "stabilityai/sd-turbo"
SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
FPS = 16
MAX_SEED = np.iinfo(np.int32).max

QUALITY_PROFILES = {
    "Schnell · 480p · 4 Schritte": {"area": 480 * 832, "steps": 4, "seconds": 70},
    "Ausgewogen · 480p · 8 Schritte": {"area": 480 * 832, "steps": 8, "seconds": 110},
    "Qualität · 576p · 12 Schritte": {"area": 576 * 1024, "steps": 12, "seconds": 145},
    "Sehr hoch · 720p · 16 Schritte": {"area": 720 * 1280, "steps": 16, "seconds": 300},
    "Maximum · 720p · 24 Schritte": {"area": 720 * 1280, "steps": 24, "seconds": 450},
}

ASPECT_RATIOS = {
    "Automatisch vom Startbild": None,
    "16:9 Querformat": 16 / 9,
    "9:16 Hochformat": 9 / 16,
    "1:1 Quadrat": 1.0,
    "4:3 Querformat": 4 / 3,
    "3:4 Hochformat": 3 / 4,
}

DEFAULT_NEGATIVE_PROMPT = (
    "static, frozen, blurry, low quality, compression artifacts, text, subtitles, watermark, "
    "deformed anatomy, extra limbs, duplicated subject, incoherent motion, camera jitter"
)


print("Loading Wan 2.2 pipeline …")
wan_pipe = WanImageToVideoPipeline.from_pretrained(
    WAN_MODEL_ID,
    transformer=WanTransformer3DModel.from_pretrained(
        WAN_WEIGHTS_ID,
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    transformer_2=WanTransformer3DModel.from_pretrained(
        WAN_WEIGHTS_ID,
        subfolder="transformer_2",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
    ),
    torch_dtype=torch.bfloat16,
)
wan_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(wan_pipe.scheduler.config, shift=8.0)
wan_pipe.to("cuda")
wan_pipe.vae.enable_tiling()

quantize_(wan_pipe.text_encoder, Int8WeightOnlyConfig())
quantize_(wan_pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
quantize_(wan_pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
spaces.aoti_blocks_load(wan_pipe.transformer, "zerogpu-aoti/Wan2", variant="fp8da")
spaces.aoti_blocks_load(wan_pipe.transformer_2, "zerogpu-aoti/Wan2", variant="fp8da")

print("Loading compact text-to-image starter …")
safety_checker = StableDiffusionSafetyChecker.from_pretrained(
    SAFETY_MODEL_ID,
    torch_dtype=torch.float16,
)
feature_extractor = CLIPImageProcessor.from_pretrained(SAFETY_MODEL_ID)
image_pipe = AutoPipelineForText2Image.from_pretrained(
    IMAGE_MODEL_ID,
    torch_dtype=torch.float16,
    variant="fp16",
    safety_checker=safety_checker,
    feature_extractor=feature_extractor,
    requires_safety_checker=True,
)
image_pipe.to("cuda")
image_pipe.set_progress_bar_config(disable=True)
print("Pipelines ready.")


def _target_size(image: Image.Image | None, aspect_choice: str, area: int) -> tuple[int, int]:
    selected_ratio = ASPECT_RATIOS.get(aspect_choice)
    if selected_ratio is None:
        selected_ratio = image.width / image.height if image is not None else 16 / 9
    width = max(16, round(np.sqrt(area * selected_ratio)) // 16 * 16)
    height = max(16, round(np.sqrt(area / selected_ratio)) // 16 * 16)
    return width, height


def _cover_resize(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    image = image.convert("RGB")
    target_width, target_height = size
    scale = max(target_width / image.width, target_height / image.height)
    resized = image.resize(
        (round(image.width * scale), round(image.height * scale)),
        Image.Resampling.LANCZOS,
    )
    left = (resized.width - target_width) // 2
    top = (resized.height - target_height) // 2
    return resized.crop((left, top, left + target_width, top + target_height))


def _to_pil_frame(frame) -> Image.Image:
    """Normalize Diffusers frame outputs to an RGB PIL image."""
    if isinstance(frame, Image.Image):
        return frame.convert("RGB")
    if isinstance(frame, torch.Tensor):
        frame = frame.detach().cpu().numpy()

    array = np.asarray(frame)
    if array.ndim == 3 and array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4):
        array = np.moveaxis(array, 0, -1)
    if np.issubdtype(array.dtype, np.floating):
        if array.size and float(array.min()) < 0:
            array = (array + 1.0) / 2.0
        array = np.clip(array, 0.0, 1.0) * 255.0
    array = np.clip(array, 0, 255).astype(np.uint8)
    return Image.fromarray(array).convert("RGB")


def _frame_count(duration_seconds: float) -> int:
    desired = int(round(float(duration_seconds) * FPS))
    return int(np.clip(round((desired - 1) / 4) * 4 + 1, 17, 97))


def estimate_gpu_duration(
    start_mode,
    uploaded_image,
    image_prompt,
    motion_prompt,
    extension_prompt,
    quality,
    aspect_ratio,
    clip_duration,
    extensions,
    negative_prompt,
    seed,
    randomize_seed,
    progress=None,
):
    del uploaded_image, image_prompt, motion_prompt, extension_prompt, aspect_ratio
    del negative_prompt, seed, randomize_seed, progress
    profile = QUALITY_PROFILES.get(quality, QUALITY_PROFILES["Ausgewogen · 480p · 8 Schritte"])
    frame_factor = _frame_count(float(clip_duration)) / _frame_count(2.1)
    total = profile["seconds"] * frame_factor * (int(extensions) + 1)
    if start_mode == "Text → Startbild → Video":
        total += 20
    # xlarge is billed at 2× quota, so 1,200s maps to the 40-minute PRO allowance.
    return min(1200, max(60, round(total)))


@spaces.GPU(duration=estimate_gpu_duration, size="xlarge")
def generate_extended_video(
    start_mode,
    uploaded_image,
    image_prompt,
    motion_prompt,
    extension_prompt,
    quality,
    aspect_ratio,
    clip_duration,
    extensions,
    negative_prompt,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    if not motion_prompt or not str(motion_prompt).strip():
        raise gr.Error("Bitte beschreibe die Bewegung im Video-Prompt.")

    extensions = int(extensions)
    segment_count = extensions + 1
    profile = QUALITY_PROFILES[quality]
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    if start_mode == "Bild → Video":
        if uploaded_image is None:
            raise gr.Error("Bitte lade ein Startbild hoch.")
        initial_source = uploaded_image.convert("RGB")
        target = _target_size(initial_source, aspect_ratio, profile["area"])
        first_frame = _cover_resize(initial_source, target)
    else:
        if not image_prompt or not str(image_prompt).strip():
            raise gr.Error("Bitte gib einen Prompt für das Startbild ein.")
        target = _target_size(None, aspect_ratio, profile["area"])
        progress(0.03, desc="Startbild aus Text erzeugen …")
        first_frame = image_pipe(
            prompt=str(image_prompt).strip(),
            negative_prompt="blurry, low quality, text, watermark",
            width=target[0],
            height=target[1],
            num_inference_steps=2,
            guidance_scale=0.0,
            generator=torch.Generator(device="cuda").manual_seed(current_seed),
        ).images[0].convert("RGB")

    current_frame = first_frame
    all_frames: list[Image.Image] = []
    clip_paths: list[str] = []
    extension_text = str(extension_prompt or "").strip()
    base_prompt = str(motion_prompt).strip()
    num_frames = _frame_count(float(clip_duration))
    print(
        f"Generation settings: quality={quality}, duration={clip_duration}s, "
        f"frames={num_frames}, extensions={extensions}, size={target[0]}x{target[1]}"
    )

    last_combined_path = None
    last_archive_path = None

    for index in range(segment_count):
        fraction = 0.08 + (index / segment_count) * 0.78
        progress(fraction, desc=f"Clip {index + 1}/{segment_count} mit Wan erzeugen …")
        prompt = base_prompt
        if index > 0:
            prompt = extension_text or (
                base_prompt
                + ", continue the action naturally from the previous clip, seamless temporal continuity"
            )

        try:
            frames = wan_pipe(
                image=current_frame,
                prompt=prompt,
                negative_prompt=negative_prompt,
                height=target[1],
                width=target[0],
                num_frames=num_frames,
                guidance_scale=1.0,
                guidance_scale_2=1.0,
                num_inference_steps=profile["steps"],
                generator=torch.Generator(device="cuda").manual_seed(
                    (current_seed + index) % MAX_SEED
                ),
            ).frames[0]
        except Exception as exc:
            traceback.print_exc()
            gc.collect()
            torch.cuda.empty_cache()
            if last_combined_path is None or last_archive_path is None:
                raise gr.Error(
                    f"Schon der erste Clip konnte nicht erzeugt werden: {type(exc).__name__}: {exc}"
                ) from exc

            status = (
                f"⚠️ Verlängerung {index} ist fehlgeschlagen. Die zuvor erzeugten {index} "
                f"Clip(s) sind unten als Video und ZIP gesichert. Lade den angezeigten letzten "
                f"Frame herunter und verwende ihn mit ‚Bild → Video‘, um dort fortzusetzen. "
                f"Fehler: {type(exc).__name__}: {exc}"
            )
            yield (
                last_combined_path,
                first_frame,
                current_frame,
                last_archive_path,
                current_seed,
                status,
            )
            return

        frames = [_to_pil_frame(frame) for frame in frames]

        clip = tempfile.NamedTemporaryFile(suffix=f"-clip-{index + 1}.mp4", delete=False)
        clip.close()
        export_to_video(frames, clip.name, fps=FPS)
        clip_paths.append(clip.name)

        if index == 0:
            all_frames.extend(frames)
        else:
            all_frames.extend(frames[1:])
        current_frame = frames[-1].convert("RGB")

        progress(
            0.08 + ((index + 0.9) / segment_count) * 0.78,
            desc=f"Checkpoint {index + 1}/{segment_count} sichern …",
        )
        combined = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.mp4", delete=False)
        combined.close()
        export_to_video(all_frames, combined.name, fps=FPS)

        archive = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.zip", delete=False)
        archive.close()
        with zipfile.ZipFile(archive.name, "w", compression=zipfile.ZIP_DEFLATED) as bundle:
            for clip_index, path in enumerate(clip_paths, start=1):
                bundle.write(path, arcname=f"clip-{clip_index:02d}.mp4")
            bundle.write(combined.name, arcname="complete-video.mp4")

        last_combined_path = combined.name
        last_archive_path = archive.name
        completed = index + 1
        total_seconds = len(all_frames) / FPS
        if completed == segment_count:
            status = (
                f"✅ Fertig: {completed} Clips ({extensions} Verlängerungen), "
                f"ca. {total_seconds:.1f} Sekunden, {target[0]}×{target[1]}, Seed {current_seed}."
            )
        else:
            status = (
                f"💾 Checkpoint gesichert: {completed}/{segment_count} Clips, "
                f"ca. {total_seconds:.1f} Sekunden. Video, ZIP und letzter Frame sind bereits verfügbar."
            )

        gc.collect()
        torch.cuda.empty_cache()
        yield combined.name, first_frame, current_frame, archive.name, current_seed, status

    # The loop yields the final checkpoint as its last result.
    return


CSS = """
.gradio-container { max-width: 1180px !important; }
.hero { text-align: center; margin: 0 auto 1.2rem; }
.hero h1 { font-size: 2.35rem; margin-bottom: .15rem; }
.flow { padding: 12px 16px; border-radius: 14px; border: 1px solid var(--border-color-primary); }
"""

with gr.Blocks() as demo:
    gr.HTML(
        "<div class='hero'><h1>♾️ Wan Extender</h1>"
        "<p>Startframe erzeugen oder hochladen · Wan-Clip generieren · automatisch verlängern</p></div>"
    )
    gr.Markdown(
        "**Logik:** Bei zwei Verlängerungen werden drei Clips erzeugt. Der letzte Frame jedes Clips "
        "wird automatisch zum ersten Frame des nächsten. Mehr Verlängerungen benötigen entsprechend "
        "mehr ZeroGPU-Kontingent."
    )

    with gr.Row():
        with gr.Column(scale=5, elem_classes="flow"):
            start_mode = gr.Radio(
                ["Text → Startbild → Video", "Bild → Video"],
                value="Text → Startbild → Video",
                label="Startmodus",
            )
            image_prompt = gr.Textbox(
                label="Prompt für das Startbild",
                placeholder="A cinematic wide shot of …",
                lines=3,
            )
            uploaded_image = gr.Image(
                type="pil",
                label="Optionales Startbild (für Modus Bild → Video)",
                sources=["upload", "clipboard"],
            )
            motion_prompt = gr.Textbox(
                label="Video- und Bewegungs-Prompt",
                placeholder="The camera slowly moves forward while …",
                lines=3,
            )
            extension_prompt = gr.Textbox(
                label="Prompt für Verlängerungen (optional)",
                placeholder="Leer lassen, um den Video-Prompt natürlich fortzuführen.",
                lines=2,
            )

        with gr.Column(scale=4, elem_classes="flow"):
            quality = gr.Dropdown(
                list(QUALITY_PROFILES),
                value="Ausgewogen · 480p · 8 Schritte",
                label="Qualitätsprofil",
            )
            aspect_ratio = gr.Dropdown(
                list(ASPECT_RATIOS),
                value="16:9 Querformat",
                label="Seitenverhältnis",
            )
            clip_duration = gr.Slider(
                1.1,
                6.0,
                value=2.1,
                step=0.1,
                label="Dauer pro Clip (Sekunden)",
                info="Bis 5 Sekunden regulär; 5,1–6 Sekunden experimenteller Langclip-Modus.",
            )
            extensions = gr.Slider(
                0,
                5,
                value=1,
                step=1,
                label="Anzahl Verlängerungen",
                info="0 = ein Clip, 5 = sechs Clips insgesamt",
            )
            with gr.Accordion("Erweiterte Einstellungen", open=False):
                negative_prompt = gr.Textbox(
                    value=DEFAULT_NEGATIVE_PROMPT,
                    label="Negative Prompt",
                    lines=3,
                )
                seed = gr.Number(value=42, precision=0, label="Seed")
                randomize_seed = gr.Checkbox(value=True, label="Seed zufällig wählen")

            generate_button = gr.Button("Video erzeugen und verlängern", variant="primary", size="lg")
            status = gr.Markdown()

    output_video = gr.Video(label="Zusammenhängendes Ergebnis", autoplay=True)
    gr.Markdown(
        "Nach jedem Clip wird sofort ein Checkpoint ausgegeben. Bei einem Fehler in einer späteren "
        "Verlängerung bleiben Video, ZIP und letzter Frame erhalten."
    )
    with gr.Row():
        generated_first_frame = gr.Image(label="Tatsächlicher erster Frame", interactive=False)
        generated_last_frame = gr.Image(label="Letzter Frame der letzten Verlängerung", interactive=False)
    clips_download = gr.File(label="Aktuellen Checkpoint als ZIP herunterladen")

    generate_button.click(
        generate_extended_video,
        inputs=[
            start_mode,
            uploaded_image,
            image_prompt,
            motion_prompt,
            extension_prompt,
            quality,
            aspect_ratio,
            clip_duration,
            extensions,
            negative_prompt,
            seed,
            randomize_seed,
        ],
        outputs=[
            output_video,
            generated_first_frame,
            generated_last_frame,
            clips_download,
            seed,
            status,
        ],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True,
        theme=gr.themes.Soft(),
        css=CSS,
    )

	from __future__ import annotations

	import gc
	import os
	import random
	import tempfile
	import traceback
	import zipfile

	# ZeroGPU currently needs this compatible CUDA/PyTorch build before importing torch.
	os.system(
	'pip install --quiet --upgrade --pre --extra-index-url '
	'https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces'
	)

	import gradio as gr
	import numpy as np
	import spaces
	import torch
	from diffusers import AutoPipelineForText2Image, FlowMatchEulerDiscreteScheduler
	from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
	from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
	from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
	from diffusers.utils.export_utils import export_to_video
	from PIL import Image
	from torchao.quantization import (
	Float8DynamicActivationFloat8WeightConfig,
	Int8WeightOnlyConfig,
	quantize_,
	)
	from transformers import CLIPImageProcessor


	WAN_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
	WAN_WEIGHTS_ID = "cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers"
	IMAGE_MODEL_ID = "stabilityai/sd-turbo"
	SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
	FPS = 16
	MAX_SEED = np.iinfo(np.int32).max

	QUALITY_PROFILES = {
	"Schnell · 480p · 4 Schritte": {"area": 480 * 832, "steps": 4, "seconds": 70},
	"Ausgewogen · 480p · 8 Schritte": {"area": 480 * 832, "steps": 8, "seconds": 110},
	"Qualität · 576p · 12 Schritte": {"area": 576 * 1024, "steps": 12, "seconds": 145},
	"Sehr hoch · 720p · 16 Schritte": {"area": 720 * 1280, "steps": 16, "seconds": 300},
	"Maximum · 720p · 24 Schritte": {"area": 720 * 1280, "steps": 24, "seconds": 450},
	}

	ASPECT_RATIOS = {
	"Automatisch vom Startbild": None,
	"16:9 Querformat": 16 / 9,
	"9:16 Hochformat": 9 / 16,
	"1:1 Quadrat": 1.0,
	"4:3 Querformat": 4 / 3,
	"3:4 Hochformat": 3 / 4,
	}

	DEFAULT_NEGATIVE_PROMPT = (
	"static, frozen, blurry, low quality, compression artifacts, text, subtitles, watermark, "
	"deformed anatomy, extra limbs, duplicated subject, incoherent motion, camera jitter"
	)


	print("Loading Wan 2.2 pipeline …")
	wan_pipe = WanImageToVideoPipeline.from_pretrained(
	WAN_MODEL_ID,
	transformer=WanTransformer3DModel.from_pretrained(
	WAN_WEIGHTS_ID,
	subfolder="transformer",
	torch_dtype=torch.bfloat16,
	device_map="cuda",
	),
	transformer_2=WanTransformer3DModel.from_pretrained(
	WAN_WEIGHTS_ID,
	subfolder="transformer_2",
	torch_dtype=torch.bfloat16,
	device_map="cuda",
	),
	torch_dtype=torch.bfloat16,
	)
	wan_pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(wan_pipe.scheduler.config, shift=8.0)
	wan_pipe.to("cuda")
	wan_pipe.vae.enable_tiling()

	quantize_(wan_pipe.text_encoder, Int8WeightOnlyConfig())
	quantize_(wan_pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
	quantize_(wan_pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
	spaces.aoti_blocks_load(wan_pipe.transformer, "zerogpu-aoti/Wan2", variant="fp8da")
	spaces.aoti_blocks_load(wan_pipe.transformer_2, "zerogpu-aoti/Wan2", variant="fp8da")

	print("Loading compact text-to-image starter …")
	safety_checker = StableDiffusionSafetyChecker.from_pretrained(
	SAFETY_MODEL_ID,
	torch_dtype=torch.float16,
	)
	feature_extractor = CLIPImageProcessor.from_pretrained(SAFETY_MODEL_ID)
	image_pipe = AutoPipelineForText2Image.from_pretrained(
	IMAGE_MODEL_ID,
	torch_dtype=torch.float16,
	variant="fp16",
	safety_checker=safety_checker,
	feature_extractor=feature_extractor,
	requires_safety_checker=True,
	)
	image_pipe.to("cuda")
	image_pipe.set_progress_bar_config(disable=True)
	print("Pipelines ready.")


	def _target_size(image: Image.Image \| None, aspect_choice: str, area: int) -> tuple[int, int]:
	selected_ratio = ASPECT_RATIOS.get(aspect_choice)
	if selected_ratio is None:
	selected_ratio = image.width / image.height if image is not None else 16 / 9
	width = max(16, round(np.sqrt(area * selected_ratio)) // 16 * 16)
	height = max(16, round(np.sqrt(area / selected_ratio)) // 16 * 16)
	return width, height


	def _cover_resize(image: Image.Image, size: tuple[int, int]) -> Image.Image:
	image = image.convert("RGB")
	target_width, target_height = size
	scale = max(target_width / image.width, target_height / image.height)
	resized = image.resize(
	(round(image.width * scale), round(image.height * scale)),
	Image.Resampling.LANCZOS,
	)
	left = (resized.width - target_width) // 2
	top = (resized.height - target_height) // 2
	return resized.crop((left, top, left + target_width, top + target_height))


	def _to_pil_frame(frame) -> Image.Image:
	"""Normalize Diffusers frame outputs to an RGB PIL image."""
	if isinstance(frame, Image.Image):
	return frame.convert("RGB")
	if isinstance(frame, torch.Tensor):
	frame = frame.detach().cpu().numpy()

	array = np.asarray(frame)
	if array.ndim == 3 and array.shape[0] in (1, 3, 4) and array.shape[-1] not in (1, 3, 4):
	array = np.moveaxis(array, 0, -1)
	if np.issubdtype(array.dtype, np.floating):
	if array.size and float(array.min()) < 0:
	array = (array + 1.0) / 2.0
	array = np.clip(array, 0.0, 1.0) * 255.0
	array = np.clip(array, 0, 255).astype(np.uint8)
	return Image.fromarray(array).convert("RGB")


	def _frame_count(duration_seconds: float) -> int:
	desired = int(round(float(duration_seconds) * FPS))
	return int(np.clip(round((desired - 1) / 4) * 4 + 1, 17, 97))


	def estimate_gpu_duration(
	start_mode,
	uploaded_image,
	image_prompt,
	motion_prompt,
	extension_prompt,
	quality,
	aspect_ratio,
	clip_duration,
	extensions,
	negative_prompt,
	seed,
	randomize_seed,
	progress=None,
	):
	del uploaded_image, image_prompt, motion_prompt, extension_prompt, aspect_ratio
	del negative_prompt, seed, randomize_seed, progress
	profile = QUALITY_PROFILES.get(quality, QUALITY_PROFILES["Ausgewogen · 480p · 8 Schritte"])
	frame_factor = _frame_count(float(clip_duration)) / _frame_count(2.1)
	total = profile["seconds"] * frame_factor * (int(extensions) + 1)
	if start_mode == "Text → Startbild → Video":
	total += 20
	# xlarge is billed at 2× quota, so 1,200s maps to the 40-minute PRO allowance.
	return min(1200, max(60, round(total)))


	@spaces.GPU(duration=estimate_gpu_duration, size="xlarge")
	def generate_extended_video(
	start_mode,
	uploaded_image,
	image_prompt,
	motion_prompt,
	extension_prompt,
	quality,
	aspect_ratio,
	clip_duration,
	extensions,
	negative_prompt,
	seed,
	randomize_seed,
	progress=gr.Progress(track_tqdm=True),
	):
	if not motion_prompt or not str(motion_prompt).strip():
	raise gr.Error("Bitte beschreibe die Bewegung im Video-Prompt.")

	extensions = int(extensions)
	segment_count = extensions + 1
	profile = QUALITY_PROFILES[quality]
	current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

	if start_mode == "Bild → Video":
	if uploaded_image is None:
	raise gr.Error("Bitte lade ein Startbild hoch.")
	initial_source = uploaded_image.convert("RGB")
	target = _target_size(initial_source, aspect_ratio, profile["area"])
	first_frame = _cover_resize(initial_source, target)
	else:
	if not image_prompt or not str(image_prompt).strip():
	raise gr.Error("Bitte gib einen Prompt für das Startbild ein.")
	target = _target_size(None, aspect_ratio, profile["area"])
	progress(0.03, desc="Startbild aus Text erzeugen …")
	first_frame = image_pipe(
	prompt=str(image_prompt).strip(),
	negative_prompt="blurry, low quality, text, watermark",
	width=target[0],
	height=target[1],
	num_inference_steps=2,
	guidance_scale=0.0,
	generator=torch.Generator(device="cuda").manual_seed(current_seed),
	).images[0].convert("RGB")

	current_frame = first_frame
	all_frames: list[Image.Image] = []
	clip_paths: list[str] = []
	extension_text = str(extension_prompt or "").strip()
	base_prompt = str(motion_prompt).strip()
	num_frames = _frame_count(float(clip_duration))
	print(
	f"Generation settings: quality={quality}, duration={clip_duration}s, "
	f"frames={num_frames}, extensions={extensions}, size={target[0]}x{target[1]}"
	)

	last_combined_path = None
	last_archive_path = None

	for index in range(segment_count):
	fraction = 0.08 + (index / segment_count) * 0.78
	progress(fraction, desc=f"Clip {index + 1}/{segment_count} mit Wan erzeugen …")
	prompt = base_prompt
	if index > 0:
	prompt = extension_text or (
	base_prompt
	+ ", continue the action naturally from the previous clip, seamless temporal continuity"
	)

	try:
	frames = wan_pipe(
	image=current_frame,
	prompt=prompt,
	negative_prompt=negative_prompt,
	height=target[1],
	width=target[0],
	num_frames=num_frames,
	guidance_scale=1.0,
	guidance_scale_2=1.0,
	num_inference_steps=profile["steps"],
	generator=torch.Generator(device="cuda").manual_seed(
	(current_seed + index) % MAX_SEED
	),
	).frames[0]
	except Exception as exc:
	traceback.print_exc()
	gc.collect()
	torch.cuda.empty_cache()
	if last_combined_path is None or last_archive_path is None:
	raise gr.Error(
	f"Schon der erste Clip konnte nicht erzeugt werden: {type(exc).__name__}: {exc}"
	) from exc

	status = (
	f"⚠️ Verlängerung {index} ist fehlgeschlagen. Die zuvor erzeugten {index} "
	f"Clip(s) sind unten als Video und ZIP gesichert. Lade den angezeigten letzten "
	f"Frame herunter und verwende ihn mit ‚Bild → Video‘, um dort fortzusetzen. "
	f"Fehler: {type(exc).__name__}: {exc}"
	)
	yield (
	last_combined_path,
	first_frame,
	current_frame,
	last_archive_path,
	current_seed,
	status,
	)
	return

	frames = [_to_pil_frame(frame) for frame in frames]

	clip = tempfile.NamedTemporaryFile(suffix=f"-clip-{index + 1}.mp4", delete=False)
	clip.close()
	export_to_video(frames, clip.name, fps=FPS)
	clip_paths.append(clip.name)

	if index == 0:
	all_frames.extend(frames)
	else:
	all_frames.extend(frames[1:])
	current_frame = frames[-1].convert("RGB")

	progress(
	0.08 + ((index + 0.9) / segment_count) * 0.78,
	desc=f"Checkpoint {index + 1}/{segment_count} sichern …",
	)
	combined = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.mp4", delete=False)
	combined.close()
	export_to_video(all_frames, combined.name, fps=FPS)

	archive = tempfile.NamedTemporaryFile(suffix="-wan-checkpoint.zip", delete=False)
	archive.close()
	with zipfile.ZipFile(archive.name, "w", compression=zipfile.ZIP_DEFLATED) as bundle:
	for clip_index, path in enumerate(clip_paths, start=1):
	bundle.write(path, arcname=f"clip-{clip_index:02d}.mp4")
	bundle.write(combined.name, arcname="complete-video.mp4")

	last_combined_path = combined.name
	last_archive_path = archive.name
	completed = index + 1
	total_seconds = len(all_frames) / FPS
	if completed == segment_count:
	status = (
	f"✅ Fertig: {completed} Clips ({extensions} Verlängerungen), "
	f"ca. {total_seconds:.1f} Sekunden, {target[0]}×{target[1]}, Seed {current_seed}."
	)
	else:
	status = (
	f"💾 Checkpoint gesichert: {completed}/{segment_count} Clips, "
	f"ca. {total_seconds:.1f} Sekunden. Video, ZIP und letzter Frame sind bereits verfügbar."
	)

	gc.collect()
	torch.cuda.empty_cache()
	yield combined.name, first_frame, current_frame, archive.name, current_seed, status

	# The loop yields the final checkpoint as its last result.
	return


	CSS = """
	.gradio-container { max-width: 1180px !important; }
	.hero { text-align: center; margin: 0 auto 1.2rem; }
	.hero h1 { font-size: 2.35rem; margin-bottom: .15rem; }
	.flow { padding: 12px 16px; border-radius: 14px; border: 1px solid var(--border-color-primary); }
	"""

	with gr.Blocks() as demo:
	gr.HTML(
	"<div class='hero'><h1>♾️ Wan Extender</h1>"
	"<p>Startframe erzeugen oder hochladen · Wan-Clip generieren · automatisch verlängern</p></div>"
	)
	gr.Markdown(
	"Logik: Bei zwei Verlängerungen werden drei Clips erzeugt. Der letzte Frame jedes Clips "
	"wird automatisch zum ersten Frame des nächsten. Mehr Verlängerungen benötigen entsprechend "
	"mehr ZeroGPU-Kontingent."
	)

	with gr.Row():
	with gr.Column(scale=5, elem_classes="flow"):
	start_mode = gr.Radio(
	["Text → Startbild → Video", "Bild → Video"],
	value="Text → Startbild → Video",
	label="Startmodus",
	)
	image_prompt = gr.Textbox(
	label="Prompt für das Startbild",
	placeholder="A cinematic wide shot of …",
	lines=3,
	)
	uploaded_image = gr.Image(
	type="pil",
	label="Optionales Startbild (für Modus Bild → Video)",
	sources=["upload", "clipboard"],
	)
	motion_prompt = gr.Textbox(
	label="Video- und Bewegungs-Prompt",
	placeholder="The camera slowly moves forward while …",
	lines=3,
	)
	extension_prompt = gr.Textbox(
	label="Prompt für Verlängerungen (optional)",
	placeholder="Leer lassen, um den Video-Prompt natürlich fortzuführen.",
	lines=2,
	)

	with gr.Column(scale=4, elem_classes="flow"):
	quality = gr.Dropdown(
	list(QUALITY_PROFILES),
	value="Ausgewogen · 480p · 8 Schritte",
	label="Qualitätsprofil",
	)
	aspect_ratio = gr.Dropdown(
	list(ASPECT_RATIOS),
	value="16:9 Querformat",
	label="Seitenverhältnis",
	)
	clip_duration = gr.Slider(
	1.1,
	6.0,
	value=2.1,
	step=0.1,
	label="Dauer pro Clip (Sekunden)",
	info="Bis 5 Sekunden regulär; 5,1–6 Sekunden experimenteller Langclip-Modus.",
	)
	extensions = gr.Slider(
	0,
	5,
	value=1,
	step=1,
	label="Anzahl Verlängerungen",
	info="0 = ein Clip, 5 = sechs Clips insgesamt",
	)
	with gr.Accordion("Erweiterte Einstellungen", open=False):
	negative_prompt = gr.Textbox(
	value=DEFAULT_NEGATIVE_PROMPT,
	label="Negative Prompt",
	lines=3,
	)
	seed = gr.Number(value=42, precision=0, label="Seed")
	randomize_seed = gr.Checkbox(value=True, label="Seed zufällig wählen")

	generate_button = gr.Button("Video erzeugen und verlängern", variant="primary", size="lg")
	status = gr.Markdown()

	output_video = gr.Video(label="Zusammenhängendes Ergebnis", autoplay=True)
	gr.Markdown(
	"Nach jedem Clip wird sofort ein Checkpoint ausgegeben. Bei einem Fehler in einer späteren "
	"Verlängerung bleiben Video, ZIP und letzter Frame erhalten."
	)
	with gr.Row():
	generated_first_frame = gr.Image(label="Tatsächlicher erster Frame", interactive=False)
	generated_last_frame = gr.Image(label="Letzter Frame der letzten Verlängerung", interactive=False)
	clips_download = gr.File(label="Aktuellen Checkpoint als ZIP herunterladen")

	generate_button.click(
	generate_extended_video,
	inputs=[
	start_mode,
	uploaded_image,
	image_prompt,
	motion_prompt,
	extension_prompt,
	quality,
	aspect_ratio,
	clip_duration,
	extensions,
	negative_prompt,
	seed,
	randomize_seed,
	],
	outputs=[
	output_video,
	generated_first_frame,
	generated_last_frame,
	clips_download,
	seed,
	status,
	],
	)

	if __name__ == "__main__":
	demo.queue(default_concurrency_limit=1).launch(
	show_error=True,
	theme=gr.themes.Soft(),
	css=CSS,
	)