Spaces:

finerobotics
/

annotation-studio

Running

Add Pika H3 Flex video generation on HF Jobs

3c65de3 verified 10 days ago

13 kB

"""Robot → human demonstration videos, one per annotated subtask.
Uses LeRobot PR #4809's ``HumanVideoModule`` from the pinned runtime bundle,
with a Studio adapter for billing and media display:
1. the VLM reads frames of the robot segment and writes the human task and an image-edit prompt;
2. Qwen-Image-Edit turns the segment's first full-resolution frame into a first-person human view;
3. the VLM writes a step-by-step video prompt from the edited frame;
4. MiniMax-H3 (480P) animates it.
Media requests go through Hugging Face Inference Providers (fal) with the signed-in user's token and
reserve their quoted price from the job's spending target before they start.
"""

import io
import json
import math
import os
import subprocess
import tempfile
import threading
import time
from pathlib import Path
from urllib.parse import urlparse

from model_catalog import human_video_profile
from PIL import Image
from provider import BudgetExceededError, DeadlineExceededError

# A video takes about three minutes; do not start one the runtime limit would cut off.
VIDEO_SECONDS_NEEDED = 420


def png_bytes(image):
    buffer = io.BytesIO()
    image.convert("RGB").save(buffer, format="PNG")
    return buffer.getvalue()


class FalMedia:
    """Qwen-Image-Edit + MiniMax-H3 on fal through the HF router, charged to ``budget``.
    ``deadline`` is the epoch by which the job must finish (None = no limit).
    """

    def __init__(self, token, budget, deadline=None, profile=None, concurrency=None, poll_interval=2.0):
        self.token, self.budget, self.deadline = token, budget, deadline
        # Pinned manifests may predate newer price fields; fill them from the current catalog.
        self.profile = {**human_video_profile(), **(profile or {})}
        self.poll_interval = poll_interval
        # Job-wide limit on simultaneous media requests, shared by every episode thread.
        self.concurrency = max(1, min(int(concurrency or self.profile["concurrency"]), 32))
        self.semaphore = threading.BoundedSemaphore(self.concurrency)

    def time_left(self):
        return self.deadline - time.time() if self.deadline else float("inf")

    def check_time(self):
        """Raise before paying for a segment that could not finish before the deadline."""
        if self.time_left() < VIDEO_SECONDS_NEEDED:
            raise DeadlineExceededError("Not enough runtime left for another video.")

    def edit_image(self, image, prompt):
        from huggingface_hub import InferenceClient

        self.check_time()
        with self.semaphore:
            self.check_time()
            reservation = self.budget.reserve_media(self.profile["edit_usd"])
            client_options = {"provider": self.profile["provider"], "api_key": self.token}
            if self.profile.get("backend") == "hf-jobs":
                # The paid GPU is idle while fal edits; bound a stalled provider request.
                client_options["timeout"] = min(300, max(1, self.time_left() - 5))
            client = InferenceClient(**client_options)
            edited = client.image_to_image(
                png_bytes(image),
                prompt=prompt,
                model=self.profile["edit_model"],
                guidance_scale=self.profile["edit_guidance_scale"],
            )
            megapixels = edited.width * edited.height / 1e6
            self.budget.record_media(reservation, megapixels * self.profile["edit_usd_per_megapixel"])
        # fal returns a square image; keep the camera's size and aspect ratio.
        return edited.convert("RGB").resize(image.size, Image.LANCZOS)

    def image_to_video(self, image, prompt):
        from huggingface_hub.inference._providers.fal_ai import FalAIImageToVideoTask
        from huggingface_hub.utils import get_session, hf_raise_for_status

        self.check_time()
        # InferenceClient has no image-text-to-video method yet; fal's image-to-video helper speaks
        # the same queue protocol once its task matches the model's mapping.
        helper = FalAIImageToVideoTask()
        helper.task = self.profile["video_task"]
        params = {
            "prompt": prompt,
            "resolution": self.profile["resolution"],
            "duration": self.profile["duration_s"],
        }
        session = get_session()
        with self.semaphore:
            request = helper.prepare_request(
                inputs=png_bytes(image),
                parameters=params,
                headers={},
                model=self.profile["video_model"],
                api_key=self.token,
            )
            reservation = self.budget.reserve_media(self.profile["video_usd"])
            response = session.post(request.url, json=request.json, headers=request.headers, timeout=60)
            hf_raise_for_status(response)
            queued = response.json()
            if not queued.get("request_id") or not queued.get("response_url"):
                raise RuntimeError("The video provider did not queue the request.")
            parsed = urlparse(request.url)
            base = f"{parsed.scheme}://{parsed.netloc}"
            if parsed.netloc == "router.huggingface.co":
                base += "/fal-ai"
            query = f"?{parsed.query}" if parsed.query else ""
            path = urlparse(queued["response_url"]).path
            status = queued.get("status")
            while status != "COMPLETED":
                if status in ("FAILED", "ERROR", "CANCELLED"):
                    raise RuntimeError(f"The video provider reported {status}.")
                if self.time_left() <= 5:
                    raise DeadlineExceededError("Runtime limit reached while a video was generating.")
                time.sleep(self.poll_interval)
                polled = session.get(f"{base}{path}/status{query}", headers=request.headers, timeout=30)
                hf_raise_for_status(polled)
                status = polled.json().get("status")
            result = session.get(f"{base}{path}{query}", headers=request.headers, timeout=60)
            hf_raise_for_status(result)
            url = result.json()["video"]["url"]
            # The result is a public CDN link; the token is never sent there.
            video = session.get(url, timeout=120)
            hf_raise_for_status(video)
            self.budget.record_media(
                reservation, self.profile["video_usd_per_second"] * self.profile["duration_s"]
            )
        return video.content


class PikaMedia(FalMedia):
    """Fal image edits and Pika's pinned open weights on the controller's A100 Job."""

    def __init__(self, token, budget, deadline=None, profile=None, concurrency=None, poll_interval=2.0):
        super().__init__(token, budget, deadline, profile, concurrency=1, poll_interval=poll_interval)

    def image_to_video(self, image, prompt):
        # Pika supports one request per GPU. Compute is already reserved in the Job quote.
        with self.semaphore:
            self.check_time()
            ratios = [(21, 9), (16, 9), (3, 2), (4, 3), (5, 4), (1, 1),
                      (4, 5), (3, 4), (2, 3), (9, 16), (9, 21)]
            ratio = image.width / image.height
            w, h = min(ratios, key=lambda pair: abs(math.log(ratio / (pair[0] / pair[1]))))
            with tempfile.TemporaryDirectory(prefix="pika-video-") as directory:
                root = Path(directory)
                (root / "first.png").write_bytes(png_bytes(image))
                (root / "prompt.txt").write_text(prompt, encoding="utf-8")
                output = root / "video.mp4"
                command = [
                    "/opt/pika-h3/bin/pika-h3-pro", "--task", "i2va",
                    "--first-image", str(root / "first.png"), "--prompt-file", str(root / "prompt.txt"),
                    "--model-repo", "Pika-Labs/Pika-H3-Flex", "--revision",
                    "0ddf55ad7067ca1e76812253bd7b4600df87f759", "--resolution", "480p",
                    "--aspect", f"{w}:{h}", "--duration", "5", "--seed", "42", "--output", str(output),
                ]
                env = {**os.environ, "HF_HUB_DISABLE_XET": "0"}
                try:
                    remaining = self.time_left()
                    timeout = max(1, remaining - 5) if math.isfinite(remaining) else None
                    subprocess.run(command, check=True, env=env, timeout=timeout)
                except subprocess.TimeoutExpired as exc:
                    raise DeadlineExceededError("Runtime limit reached while a Pika video was generating.") from exc
                if not output.is_file() or not output.stat().st_size:
                    raise RuntimeError("Pika did not produce a video.")
                return output.read_bytes()


def generate_human_videos(record, staging, vlm, frames, full_frames, media, out_dir, max_segments):
    """Use PR #4809's module with Studio billing, full-resolution edits and UI entries."""
    from lerobot.annotations.steerable_pipeline.config import HumanVideoConfig
    from lerobot.annotations.steerable_pipeline.modules.human_video import HumanVideoModule

    profile = media.profile
    stop = threading.Event()
    root = out_dir.parent.parent

    class FrameAdapter:
        camera_keys = getattr(frames, "camera_keys", [])
        camera_key = getattr(frames, "camera_key", None)

        def frames_at(self, record, timestamps, camera_key=None):
            # The core module requests first-frame + six context frames together.
            return [
                *full_frames.frames_at(record, timestamps[:1], camera_key=camera_key),
                *frames.frames_at(record, timestamps[1:], camera_key=camera_key),
            ]

    class GeneratorAdapter:
        def edit_image(self, image, prompt):
            return media.edit_image(image, prompt)

        def image_to_video(self, image, prompt, end_image=None):
            if end_image is not None:
                raise ValueError("The Studio quote covers a single edited first frame")
            return media.image_to_video(image, prompt)

    class StudioHumanVideo(HumanVideoModule):
        def _generate_segment(self, record, k, span, context, plan, camera, directory):
            entry = {
                "source_episode_index": record.episode_index,
                "segment_index": k,
                "subtask": span["text"],
                "start_timestamp": span["start"],
                "end_timestamp": span["end"],
                "files": {},
            }
            if stop.is_set():
                return {**entry, "status": "skipped"}
            try:
                media.check_time()
            except BudgetExceededError as exc:
                stop.set()
                return {
                    **entry,
                    "status": "time_stopped" if isinstance(exc, DeadlineExceededError) else "budget_stopped",
                }
            robot_path = directory / f"segment_{k:03d}_robot_frame.png"
            context[0].convert("RGB").save(robot_path)
            generated = super()._generate_segment(record, k, span, context, plan, camera, directory)
            entry.update(generated)
            entry.pop("episode_index", None)
            entry["status"] = "completed" if generated["status"] == "ok" else generated["status"]
            if generated.get("error_type") in {"BudgetExceededError", "DeadlineExceededError"}:
                stop.set()
                entry["status"] = (
                    "time_stopped" if generated["error_type"] == "DeadlineExceededError" else "budget_stopped"
                )
            entry["files"] = {"robot_frame": str(robot_path)}
            for kind in ("first_frame", "video"):
                path = generated.get(kind + "_path")
                if path:
                    entry["files"][kind] = str(root / path)
                    entry.pop(kind + "_path", None)
            # Failed segments may still have a successfully edited first frame.
            first = directory / f"segment_{k:03d}_first_frame.png"
            if first.exists():
                entry["files"]["first_frame"] = str(first)
            return entry

    cfg = HumanVideoConfig(
        enabled=True,
        provider="hf-jobs" if profile.get("backend") == "hf-jobs" else profile["provider"],
        edit_model=profile["edit_model"],
        edit_guidance_scale=profile.get("edit_guidance_scale", 1.5),
        video_model=profile["video_model"],
        resolution=profile["resolution"],
        duration_s=profile["duration_s"],
        end_frame=False,
        max_segments_per_episode=max_segments or None,
        max_concurrency=max(1, getattr(media, "concurrency", 1)),
        output_dir=out_dir.parent.name,
    )
    module = StudioHumanVideo(
        vlm=vlm, config=cfg, root=root, generator=GeneratorAdapter(), frame_provider=FrameAdapter()
    )
    module.run_episode(record, staging)
    return json.loads((staging.episode_dir / "human_video.json").read_text())

	"""Robot → human demonstration videos, one per annotated subtask.

	Uses LeRobot PR #4809's ``HumanVideoModule`` from the pinned runtime bundle,
	with a Studio adapter for billing and media display:

	1. the VLM reads frames of the robot segment and writes the human task and an image-edit prompt;
	2. Qwen-Image-Edit turns the segment's first full-resolution frame into a first-person human view;
	3. the VLM writes a step-by-step video prompt from the edited frame;
	4. MiniMax-H3 (480P) animates it.

	Media requests go through Hugging Face Inference Providers (fal) with the signed-in user's token and
	reserve their quoted price from the job's spending target before they start.
	"""

	import io
	import json
	import math
	import os
	import subprocess
	import tempfile
	import threading
	import time
	from pathlib import Path
	from urllib.parse import urlparse

	from model_catalog import human_video_profile
	from PIL import Image
	from provider import BudgetExceededError, DeadlineExceededError

	# A video takes about three minutes; do not start one the runtime limit would cut off.
	VIDEO_SECONDS_NEEDED = 420


	def png_bytes(image):
	buffer = io.BytesIO()
	image.convert("RGB").save(buffer, format="PNG")
	return buffer.getvalue()


	class FalMedia:
	"""Qwen-Image-Edit + MiniMax-H3 on fal through the HF router, charged to ``budget``.

	``deadline`` is the epoch by which the job must finish (None = no limit).
	"""

	def __init__(self, token, budget, deadline=None, profile=None, concurrency=None, poll_interval=2.0):
	self.token, self.budget, self.deadline = token, budget, deadline
	# Pinned manifests may predate newer price fields; fill them from the current catalog.
	self.profile = {human_video_profile(), (profile or {})}
	self.poll_interval = poll_interval
	# Job-wide limit on simultaneous media requests, shared by every episode thread.
	self.concurrency = max(1, min(int(concurrency or self.profile["concurrency"]), 32))
	self.semaphore = threading.BoundedSemaphore(self.concurrency)

	def time_left(self):
	return self.deadline - time.time() if self.deadline else float("inf")

	def check_time(self):
	"""Raise before paying for a segment that could not finish before the deadline."""
	if self.time_left() < VIDEO_SECONDS_NEEDED:
	raise DeadlineExceededError("Not enough runtime left for another video.")

	def edit_image(self, image, prompt):
	from huggingface_hub import InferenceClient

	self.check_time()
	with self.semaphore:
	self.check_time()
	reservation = self.budget.reserve_media(self.profile["edit_usd"])
	client_options = {"provider": self.profile["provider"], "api_key": self.token}
	if self.profile.get("backend") == "hf-jobs":
	# The paid GPU is idle while fal edits; bound a stalled provider request.
	client_options["timeout"] = min(300, max(1, self.time_left() - 5))
	client = InferenceClient(**client_options)
	edited = client.image_to_image(
	png_bytes(image),
	prompt=prompt,
	model=self.profile["edit_model"],
	guidance_scale=self.profile["edit_guidance_scale"],
	)
	megapixels = edited.width * edited.height / 1e6
	self.budget.record_media(reservation, megapixels * self.profile["edit_usd_per_megapixel"])
	# fal returns a square image; keep the camera's size and aspect ratio.
	return edited.convert("RGB").resize(image.size, Image.LANCZOS)

	def image_to_video(self, image, prompt):
	from huggingface_hub.inference._providers.fal_ai import FalAIImageToVideoTask
	from huggingface_hub.utils import get_session, hf_raise_for_status

	self.check_time()
	# InferenceClient has no image-text-to-video method yet; fal's image-to-video helper speaks
	# the same queue protocol once its task matches the model's mapping.
	helper = FalAIImageToVideoTask()
	helper.task = self.profile["video_task"]
	params = {
	"prompt": prompt,
	"resolution": self.profile["resolution"],
	"duration": self.profile["duration_s"],
	}
	session = get_session()
	with self.semaphore:
	request = helper.prepare_request(
	inputs=png_bytes(image),
	parameters=params,
	headers={},
	model=self.profile["video_model"],
	api_key=self.token,
	)
	reservation = self.budget.reserve_media(self.profile["video_usd"])
	response = session.post(request.url, json=request.json, headers=request.headers, timeout=60)
	hf_raise_for_status(response)
	queued = response.json()
	if not queued.get("request_id") or not queued.get("response_url"):
	raise RuntimeError("The video provider did not queue the request.")
	parsed = urlparse(request.url)
	base = f"{parsed.scheme}://{parsed.netloc}"
	if parsed.netloc == "router.huggingface.co":
	base += "/fal-ai"
	query = f"?{parsed.query}" if parsed.query else ""
	path = urlparse(queued["response_url"]).path
	status = queued.get("status")
	while status != "COMPLETED":
	if status in ("FAILED", "ERROR", "CANCELLED"):
	raise RuntimeError(f"The video provider reported {status}.")
	if self.time_left() <= 5:
	raise DeadlineExceededError("Runtime limit reached while a video was generating.")
	time.sleep(self.poll_interval)
	polled = session.get(f"{base}{path}/status{query}", headers=request.headers, timeout=30)
	hf_raise_for_status(polled)
	status = polled.json().get("status")
	result = session.get(f"{base}{path}{query}", headers=request.headers, timeout=60)
	hf_raise_for_status(result)
	url = result.json()["video"]["url"]
	# The result is a public CDN link; the token is never sent there.
	video = session.get(url, timeout=120)
	hf_raise_for_status(video)
	self.budget.record_media(
	reservation, self.profile["video_usd_per_second"] * self.profile["duration_s"]
	)
	return video.content


	class PikaMedia(FalMedia):
	"""Fal image edits and Pika's pinned open weights on the controller's A100 Job."""

	def __init__(self, token, budget, deadline=None, profile=None, concurrency=None, poll_interval=2.0):
	super().__init__(token, budget, deadline, profile, concurrency=1, poll_interval=poll_interval)

	def image_to_video(self, image, prompt):
	# Pika supports one request per GPU. Compute is already reserved in the Job quote.
	with self.semaphore:
	self.check_time()
	ratios = [(21, 9), (16, 9), (3, 2), (4, 3), (5, 4), (1, 1),
	(4, 5), (3, 4), (2, 3), (9, 16), (9, 21)]
	ratio = image.width / image.height
	w, h = min(ratios, key=lambda pair: abs(math.log(ratio / (pair[0] / pair[1]))))
	with tempfile.TemporaryDirectory(prefix="pika-video-") as directory:
	root = Path(directory)
	(root / "first.png").write_bytes(png_bytes(image))
	(root / "prompt.txt").write_text(prompt, encoding="utf-8")
	output = root / "video.mp4"
	command = [
	"/opt/pika-h3/bin/pika-h3-pro", "--task", "i2va",
	"--first-image", str(root / "first.png"), "--prompt-file", str(root / "prompt.txt"),
	"--model-repo", "Pika-Labs/Pika-H3-Flex", "--revision",
	"0ddf55ad7067ca1e76812253bd7b4600df87f759", "--resolution", "480p",
	"--aspect", f"{w}:{h}", "--duration", "5", "--seed", "42", "--output", str(output),
	]
	env = {**os.environ, "HF_HUB_DISABLE_XET": "0"}
	try:
	remaining = self.time_left()
	timeout = max(1, remaining - 5) if math.isfinite(remaining) else None
	subprocess.run(command, check=True, env=env, timeout=timeout)
	except subprocess.TimeoutExpired as exc:
	raise DeadlineExceededError("Runtime limit reached while a Pika video was generating.") from exc
	if not output.is_file() or not output.stat().st_size:
	raise RuntimeError("Pika did not produce a video.")
	return output.read_bytes()


	def generate_human_videos(record, staging, vlm, frames, full_frames, media, out_dir, max_segments):
	"""Use PR #4809's module with Studio billing, full-resolution edits and UI entries."""
	from lerobot.annotations.steerable_pipeline.config import HumanVideoConfig
	from lerobot.annotations.steerable_pipeline.modules.human_video import HumanVideoModule

	profile = media.profile
	stop = threading.Event()
	root = out_dir.parent.parent

	class FrameAdapter:
	camera_keys = getattr(frames, "camera_keys", [])
	camera_key = getattr(frames, "camera_key", None)

	def frames_at(self, record, timestamps, camera_key=None):
	# The core module requests first-frame + six context frames together.
	return [
	*full_frames.frames_at(record, timestamps[:1], camera_key=camera_key),
	*frames.frames_at(record, timestamps[1:], camera_key=camera_key),
	]

	class GeneratorAdapter:
	def edit_image(self, image, prompt):
	return media.edit_image(image, prompt)

	def image_to_video(self, image, prompt, end_image=None):
	if end_image is not None:
	raise ValueError("The Studio quote covers a single edited first frame")
	return media.image_to_video(image, prompt)

	class StudioHumanVideo(HumanVideoModule):
	def _generate_segment(self, record, k, span, context, plan, camera, directory):
	entry = {
	"source_episode_index": record.episode_index,
	"segment_index": k,
	"subtask": span["text"],
	"start_timestamp": span["start"],
	"end_timestamp": span["end"],
	"files": {},
	}
	if stop.is_set():
	return {**entry, "status": "skipped"}
	try:
	media.check_time()
	except BudgetExceededError as exc:
	stop.set()
	return {
	**entry,
	"status": "time_stopped" if isinstance(exc, DeadlineExceededError) else "budget_stopped",
	}
	robot_path = directory / f"segment_{k:03d}_robot_frame.png"
	context[0].convert("RGB").save(robot_path)
	generated = super()._generate_segment(record, k, span, context, plan, camera, directory)
	entry.update(generated)
	entry.pop("episode_index", None)
	entry["status"] = "completed" if generated["status"] == "ok" else generated["status"]
	if generated.get("error_type") in {"BudgetExceededError", "DeadlineExceededError"}:
	stop.set()
	entry["status"] = (
	"time_stopped" if generated["error_type"] == "DeadlineExceededError" else "budget_stopped"
	)
	entry["files"] = {"robot_frame": str(robot_path)}
	for kind in ("first_frame", "video"):
	path = generated.get(kind + "_path")
	if path:
	entry["files"][kind] = str(root / path)
	entry.pop(kind + "_path", None)
	# Failed segments may still have a successfully edited first frame.
	first = directory / f"segment_{k:03d}_first_frame.png"
	if first.exists():
	entry["files"]["first_frame"] = str(first)
	return entry

	cfg = HumanVideoConfig(
	enabled=True,
	provider="hf-jobs" if profile.get("backend") == "hf-jobs" else profile["provider"],
	edit_model=profile["edit_model"],
	edit_guidance_scale=profile.get("edit_guidance_scale", 1.5),
	video_model=profile["video_model"],
	resolution=profile["resolution"],
	duration_s=profile["duration_s"],
	end_frame=False,
	max_segments_per_episode=max_segments or None,
	max_concurrency=max(1, getattr(media, "concurrency", 1)),
	output_dir=out_dir.parent.name,
	)
	module = StudioHumanVideo(
	vlm=vlm, config=cfg, root=root, generator=GeneratorAdapter(), frame_provider=FrameAdapter()
	)
	module.run_episode(record, staging)
	return json.loads((staging.episode_dir / "human_video.json").read_text())