File size: 6,423 Bytes
055bfe3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """HF Inference Endpoint handler for SulphurAI/Sulphur-2-base (LTX-2.3 fine-tune).
The pipeline is constructed against the low-level `ltx_pipelines.DistilledPipeline`
API rather than `diffusers`, because Sulphur ships raw .safetensors and explicitly
recommends the distilled checkpoint + distill LoRA workflow.
Required environment:
- LTX-2 (Lightricks/LTX-2) installed: `ltx_core`, `ltx_pipelines` importable.
The official install path is `uv sync` from a git clone — see Dockerfile.
- HF_TOKEN set as an Endpoint secret. The Gemma text encoder is gated, so the
HF account behind the token must have accepted its license.
- GPU with >= 48 GB VRAM (H100 80GB / A100 80GB recommended). Distilled bf16
plus Gemma 12B comfortably exceeds 40 GB at fp16.
"""
from __future__ import annotations
import base64
import logging
import os
import tempfile
import uuid
from typing import Any, Dict, Iterator
import torch
from huggingface_hub import hf_hub_download, snapshot_download
logger = logging.getLogger("sulphur2.handler")
logging.basicConfig(level=logging.INFO)
SULPHUR_REPO = "SulphurAI/Sulphur-2-base"
SULPHUR_DISTILLED = "sulphur_distil_bf16.safetensors"
SULPHUR_DISTILL_LORA = "distill_loras/ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors"
LTX_REPO = "Lightricks/LTX-2.3"
LTX_SPATIAL_UPSCALER = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
def _import_ltx():
"""Import LTX-2 lazily so import errors surface at handler init, not module load."""
from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_pipelines import DistilledPipeline
from ltx_pipelines.utils.media_io import encode_video
return DistilledPipeline, LoraPathStrengthAndSDOps, encode_video
class EndpointHandler:
def __init__(self, model_dir: str, **kwargs: Any) -> None:
token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
if not token:
raise RuntimeError(
"HF_TOKEN (or HUGGING_FACE_HUB_TOKEN) is required. Accept the Gemma "
"license on huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized "
"and add the token as an Endpoint secret."
)
DistilledPipeline, LoraPathStrengthAndSDOps, encode_video = _import_ltx()
self._encode_video = encode_video
logger.info("Downloading Sulphur distilled checkpoint...")
checkpoint_path = hf_hub_download(
repo_id=SULPHUR_REPO, filename=SULPHUR_DISTILLED, token=token
)
logger.info("Downloading Sulphur distill LoRA...")
distill_lora_path = hf_hub_download(
repo_id=SULPHUR_REPO, filename=SULPHUR_DISTILL_LORA, token=token
)
logger.info("Downloading LTX-2.3 spatial upscaler...")
spatial_upscaler_path = hf_hub_download(
repo_id=LTX_REPO, filename=LTX_SPATIAL_UPSCALER, token=token
)
logger.info("Downloading Gemma 3 text encoder snapshot...")
gemma_root = snapshot_download(
repo_id=GEMMA_REPO,
token=token,
allow_patterns=[
"*.json",
"*.model",
"*.safetensors",
"tokenizer*",
"special_tokens_map.json",
],
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info("Initializing DistilledPipeline on %s...", device)
loras = [LoraPathStrengthAndSDOps(path=distill_lora_path, strength=1.0)]
self.pipeline = DistilledPipeline(
distilled_checkpoint_path=checkpoint_path,
gemma_root=gemma_root,
spatial_upsampler_path=spatial_upscaler_path,
loras=loras,
device=device,
)
logger.info("Pipeline ready.")
@torch.inference_mode()
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
prompt = data.get("inputs") or data.get("prompt")
if not isinstance(prompt, str) or not prompt.strip():
return {"error": "Request must include a non-empty 'inputs' or 'prompt' string."}
params = data.get("parameters") or {}
width = int(params.get("width", 768))
height = int(params.get("height", 512))
num_frames = int(params.get("num_frames", 97))
fps = int(params.get("fps", 24))
seed = int(params.get("seed", 0))
enhance_prompt = bool(params.get("enhance_prompt", False))
if num_frames < 9 or num_frames > 257:
return {"error": "num_frames must be between 9 and 257."}
if width % 32 or height % 32:
return {"error": "width and height must be multiples of 32."}
logger.info(
"Generating: prompt=%r seed=%d %dx%d frames=%d fps=%d",
prompt[:80], seed, width, height, num_frames, fps,
)
video_iter, audio = self.pipeline(
prompt=prompt,
seed=seed,
height=height,
width=width,
num_frames=num_frames,
frame_rate=float(fps),
images=[],
tiling_config=None,
enhance_prompt=enhance_prompt,
)
chunks = list(video_iter) if isinstance(video_iter, Iterator) else [video_iter]
if not chunks:
return {"error": "Pipeline produced no frames."}
video_chunks_number = len(chunks)
tmp_path = os.path.join(tempfile.gettempdir(), f"sulphur_{uuid.uuid4().hex}.mp4")
try:
self._encode_video(
video=iter(chunks),
fps=fps,
audio=audio,
output_path=tmp_path,
video_chunks_number=video_chunks_number,
)
with open(tmp_path, "rb") as f:
payload = f.read()
finally:
if os.path.exists(tmp_path):
try:
os.remove(tmp_path)
except OSError:
pass
return {
"content_type": "video/mp4",
"encoding": "base64",
"video_base64": base64.b64encode(payload).decode("ascii"),
"metadata": {
"width": width,
"height": height,
"num_frames": num_frames,
"fps": fps,
"seed": seed,
},
}
|