Spaces:
Running on Zero
Running on Zero
Upload folder using huggingface_hub
Browse files
app.py
CHANGED
|
@@ -16,6 +16,7 @@ encode of the structure video and first frame β runs here, on the card, inside
|
|
| 16 |
|
| 17 |
from __future__ import annotations
|
| 18 |
|
|
|
|
| 19 |
import json
|
| 20 |
import os
|
| 21 |
import sys
|
|
@@ -25,6 +26,15 @@ import traceback
|
|
| 25 |
from functools import cache
|
| 26 |
from pathlib import Path
|
| 27 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 28 |
# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the model
|
| 29 |
# load can happen at startup rather than on GPU time.
|
| 30 |
import spaces
|
|
@@ -151,6 +161,22 @@ def load_models() -> str | None:
|
|
| 151 |
return LOAD_ERROR
|
| 152 |
|
| 153 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 154 |
# ββ the encode chain, over the one VAE ββββββββββββββββββββββββββββββββββββββ
|
| 155 |
|
| 156 |
_CHAIN = {}
|
|
@@ -350,11 +376,16 @@ def _generate(structure, first, embeds, tags, height, width, seed):
|
|
| 350 |
if PLACEMENT == "lazy" and CORE is not None:
|
| 351 |
CORE.model.to("cuda")
|
| 352 |
CORE.device = torch.device("cuda")
|
| 353 |
-
VAE.to("cuda")
|
| 354 |
|
| 355 |
chain = latents_chain()
|
| 356 |
from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3VideoReference
|
| 357 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 358 |
# The references exactly as the release's `encode()` built them: the structure video
|
| 359 |
# center-cropped to the canvas and prepared, the first frame as a STILL_FRAMES-long still.
|
| 360 |
video, fps = pre.read_video(structure, FRAMES / pre.FPS + 0.5)
|
|
@@ -409,19 +440,34 @@ def _generate(structure, first, embeds, tags, height, width, seed):
|
|
| 409 |
{"observations": observations, "latent_frames": frame_count, "canvas": [height, width]}
|
| 410 |
)})
|
| 411 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 412 |
stream = object.__new__(H3OmniStream)
|
| 413 |
stream.root, stream.current_plan_fingerprint, stream.aligned_control = "/", "render", False
|
| 414 |
batch = stream.read({"id": "request", "member": member})
|
| 415 |
began = time.time()
|
|
|
|
| 416 |
with torch.no_grad():
|
| 417 |
latents = CORE.generate(batch, noise_seed=int(seed), num_inference_steps=FORWARDS + 1, renoise=True)
|
| 418 |
sampled = time.time() - began
|
|
|
|
|
|
|
| 419 |
decoded = decode_h3_latents(VAE, latents, device=CORE.device)
|
| 420 |
data = _mp4_bytes(decoded)
|
|
|
|
| 421 |
os.unlink(member)
|
|
|
|
|
|
|
| 422 |
print(
|
| 423 |
f"[{VERSION}] sampled {FORWARDS} steps in {sampled:.0f}s; "
|
| 424 |
-
f"{
|
| 425 |
flush=True,
|
| 426 |
)
|
| 427 |
return data
|
|
@@ -513,6 +559,29 @@ CSS = """
|
|
| 513 |
.dark .gradio-container { color: var(--body-text-color); }
|
| 514 |
"""
|
| 515 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 516 |
with gr.Blocks(title="SolarRenderer") as demo:
|
| 517 |
gr.Markdown(INTRO)
|
| 518 |
|
|
@@ -537,7 +606,7 @@ with gr.Blocks(title="SolarRenderer") as demo:
|
|
| 537 |
gr.Examples(
|
| 538 |
# The release's own example: the semantic-map highway structure, its first frame, and the
|
| 539 |
# full synthwave prompt the authors ship in examples/sem-highway_synthwave/.
|
| 540 |
-
examples=[[
|
| 541 |
inputs=[prompt, structure, first],
|
| 542 |
outputs=result,
|
| 543 |
fn=generate,
|
|
|
|
| 16 |
|
| 17 |
from __future__ import annotations
|
| 18 |
|
| 19 |
+
import gc
|
| 20 |
import json
|
| 21 |
import os
|
| 22 |
import sys
|
|
|
|
| 26 |
from functools import cache
|
| 27 |
from pathlib import Path
|
| 28 |
|
| 29 |
+
# The render packs ~96k rows through a 50-layer 33B trunk: the feed-forward alone allocates ~5.6 GB
|
| 30 |
+
# in one tensor, and the allocator's default 2 MB-block caches cannot hand back a run that large
|
| 31 |
+
# once the VAE encode has left a few GiB of odd-sized blocks behind β the failure reads as
|
| 32 |
+
# `NVML_SUCCESS == r INTERNAL ASSERT FAILED ... CUDACachingAllocator.cpp` with free memory still
|
| 33 |
+
# showing on the card. Expandable segments let one segment grow instead of needing contiguity.
|
| 34 |
+
_ALLOC = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
|
| 35 |
+
if "expandable_segments" not in _ALLOC:
|
| 36 |
+
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = f"{_ALLOC},".lstrip(",") + "expandable_segments:True"
|
| 37 |
+
|
| 38 |
# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the model
|
| 39 |
# load can happen at startup rather than on GPU time.
|
| 40 |
import spaces
|
|
|
|
| 161 |
return LOAD_ERROR
|
| 162 |
|
| 163 |
|
| 164 |
+
def release_card() -> None:
|
| 165 |
+
"""Send the conditioners home and give their blocks back to the driver.
|
| 166 |
+
|
| 167 |
+
The release runs the encoder and the renderer in separate processes because they do not fit
|
| 168 |
+
together: `del encoder; free_cuda()`. Here they share a process, so the same release has to be
|
| 169 |
+
spelled out β the float32 visual VAE (9.7 GiB) plus the audio VAE (0.56 GiB) plus whatever the
|
| 170 |
+
158-frame pixel encode left in the allocator is exactly the headroom the 33B trunk needs for the
|
| 171 |
+
~5.6 GB feed-forward activation each of the four evaluations allocates in one tensor.
|
| 172 |
+
"""
|
| 173 |
+
for module in (VAE, getattr(CONDITIONERS, "audio_vae", None)):
|
| 174 |
+
if isinstance(module, torch.nn.Module):
|
| 175 |
+
module.to("cpu")
|
| 176 |
+
gc.collect()
|
| 177 |
+
torch.cuda.empty_cache()
|
| 178 |
+
|
| 179 |
+
|
| 180 |
# ββ the encode chain, over the one VAE ββββββββββββββββββββββββββββββββββββββ
|
| 181 |
|
| 182 |
_CHAIN = {}
|
|
|
|
| 376 |
if PLACEMENT == "lazy" and CORE is not None:
|
| 377 |
CORE.model.to("cuda")
|
| 378 |
CORE.device = torch.device("cuda")
|
|
|
|
| 379 |
|
| 380 |
chain = latents_chain()
|
| 381 |
from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3VideoReference
|
| 382 |
|
| 383 |
+
# The structure encode below needs the visual VAE on the card; the sample that follows does not
|
| 384 |
+
# (see the release below). Bring back whatever an earlier request sent home.
|
| 385 |
+
VAE.to(CORE.device)
|
| 386 |
+
if getattr(CONDITIONERS, "audio_vae", None) is not None:
|
| 387 |
+
CONDITIONERS.audio_vae.to(CORE.device)
|
| 388 |
+
|
| 389 |
# The references exactly as the release's `encode()` built them: the structure video
|
| 390 |
# center-cropped to the canvas and prepared, the first frame as a STILL_FRAMES-long still.
|
| 391 |
video, fps = pre.read_video(structure, FRAMES / pre.FPS + 0.5)
|
|
|
|
| 440 |
{"observations": observations, "latent_frames": frame_count, "canvas": [height, width]}
|
| 441 |
)})
|
| 442 |
|
| 443 |
+
# Everything the encode needed is now in the member file, so hand the card back before sampling.
|
| 444 |
+
# The release does the same thing with `del encoder; free_cuda()` β its encoder and its renderer
|
| 445 |
+
# never share the card β and here it is what makes room for the ~5.6 GB feed-forward activation
|
| 446 |
+
# the 33B trunk allocates per step. The VAE is 9.7 GiB of float32 (the release's own validation
|
| 447 |
+
# decoder runs in FP32) and the 158-frame pixel encode leaves tensors of the same order behind,
|
| 448 |
+
# so both it and the reference stack go home first; the decoder gets them back below.
|
| 449 |
+
del state, target_state, refs, structure_frames, target, tensors
|
| 450 |
+
release_card()
|
| 451 |
+
|
| 452 |
stream = object.__new__(H3OmniStream)
|
| 453 |
stream.root, stream.current_plan_fingerprint, stream.aligned_control = "/", "render", False
|
| 454 |
batch = stream.read({"id": "request", "member": member})
|
| 455 |
began = time.time()
|
| 456 |
+
# `no_grad`, as the release's own `render_member` samples.
|
| 457 |
with torch.no_grad():
|
| 458 |
latents = CORE.generate(batch, noise_seed=int(seed), num_inference_steps=FORWARDS + 1, renoise=True)
|
| 459 |
sampled = time.time() - began
|
| 460 |
+
del batch
|
| 461 |
+
VAE.to(CORE.device) # the decode needs the visual VAE, and only the decode
|
| 462 |
decoded = decode_h3_latents(VAE, latents, device=CORE.device)
|
| 463 |
data = _mp4_bytes(decoded)
|
| 464 |
+
shape = tuple(decoded.shape)
|
| 465 |
os.unlink(member)
|
| 466 |
+
del decoded, latents
|
| 467 |
+
release_card()
|
| 468 |
print(
|
| 469 |
f"[{VERSION}] sampled {FORWARDS} steps in {sampled:.0f}s; "
|
| 470 |
+
f"{shape[-3]} frames at {shape[-2]}x{shape[-1]}",
|
| 471 |
flush=True,
|
| 472 |
)
|
| 473 |
return data
|
|
|
|
| 559 |
.dark .gradio-container { color: var(--body-text-color); }
|
| 560 |
"""
|
| 561 |
|
| 562 |
+
# The release's own `examples/sem-highway_synthwave/prompt.txt`, inlined: a `gr.Textbox` example
|
| 563 |
+
# cell is a *string*, not a path β Gradio would paste `examples/prompt.txt` into the box verbatim
|
| 564 |
+
# rather than reading the file.
|
| 565 |
+
EXAMPLE_PROMPT = """subject_definitions:
|
| 566 |
+
<Video 1> is a semantic segmentation of the scene, one flat colour per class (road, sidewalk, building, vegetation, vehicle, pedestrian, sky and so on), frame by frame, from the same camera as the target video.
|
| 567 |
+
<Video 2> is the first frame of [Shot 1], held still.
|
| 568 |
+
|
| 569 |
+
summary:
|
| 570 |
+
[keyframe completion + reference generation] The target video begins from <Video 2> and follows the frame-by-frame structure and motion of <Video 1>. A retro synthwave highway run: a neon-trimmed coupe overtakes a glowing semi-truck, flies under a neon gantry and overpass at dusk, and slides back into its lane.
|
| 571 |
+
|
| 572 |
+
retention_analysis:
|
| 573 |
+
<Video 1> (frame-by-frame layout of every class and camera): attribute_transfer - every region keeps its class, outline and place; materials, colours and lighting are not given by <Video 1>; the flat class colours are not shown.
|
| 574 |
+
<Video 2> ([Shot 1] first frame): fully_preserved - the composition, subjects, materials, colours and lighting of <Video 2> are the target video's first frame.
|
| 575 |
+
|
| 576 |
+
detailed_description:
|
| 577 |
+
The target video is a stylised synthwave animation in the look of an eighties retro-futuristic music video, with a magenta-to-violet sunset gradient sky, a huge striped setting sun on the horizon, neon pink and cyan edge lighting, glossy reflective surfaces and a light film grain. [Shot 1] The shot begins from <Video 2>. The camera trails a few metres behind and above a black wedge-shaped sports coupe with a glowing cyan light strip across its tail, in the right lane of a six-lane divided highway whose lane lines glow hot pink. Ahead sits the back of a chrome semi-trailer outlined in purple neon, pulled by a tractor with a glowing grille. After half a second the coupe glides left into the middle lane, the camera tilting with it, and begins to close on the truck. The trailer's mirror-like side slides by on the right, reflecting the sunset, then its glowing wheels and the cab with chrome stacks, passing within a couple of metres. Midway, a gantry with two glowing signs in retro lettering sweeps overhead. Further up the road, tail lights of an SUV and a sedan streak away in the fast lane, with a hatchback, a box truck and a van glowing ahead. Across a dark median lined with tall lamp posts topped by pink light bars, oncoming cars and a coach race past as streaks of white light. Near the end the coupe flashes under a concrete overpass trimmed with neon tubes, while a glowing SUV and a delivery truck cross above, their light trails painting streaks along the bridge railing. Clear of the truck, the coupe slides back into the right lane toward the setting sun. The coupe stays planted through every move, rocking slightly over expansion joints, and the camera follows with a slight lag so each lane change feels smooth, heavy and rhythmic, like a slow camera move in a music video. Guard rails gleam with reflected pink light, palm-like trees and stylised pines stand as black silhouettes, utility poles and their wires cut thin lines across the gradient sky, and low warehouses with glowing windows and purple hills fade toward a wireframe-like horizon.
|
| 578 |
+
|
| 579 |
+
overall_soundscape:
|
| 580 |
+
Smooth tyre roar and wind with a punchy engine note as the coupe overtakes, a deep whoosh past the truck and a hollow echo under the overpass.
|
| 581 |
+
|
| 582 |
+
non_diegetic_music:
|
| 583 |
+
N/A"""
|
| 584 |
+
|
| 585 |
with gr.Blocks(title="SolarRenderer") as demo:
|
| 586 |
gr.Markdown(INTRO)
|
| 587 |
|
|
|
|
| 606 |
gr.Examples(
|
| 607 |
# The release's own example: the semantic-map highway structure, its first frame, and the
|
| 608 |
# full synthwave prompt the authors ship in examples/sem-highway_synthwave/.
|
| 609 |
+
examples=[[EXAMPLE_PROMPT, "examples/structure.mp4", "examples/first.png"]],
|
| 610 |
inputs=[prompt, structure, first],
|
| 611 |
outputs=result,
|
| 612 |
fn=generate,
|