multimodalart HF Staff commited on
Commit
59dca01
Β·
verified Β·
1 Parent(s): 2acaa72

Upload folder using huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +72 -3
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"{decoded.shape[-3]} frames at {decoded.shape[-2]}x{decoded.shape[-1]}",
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=[["examples/prompt.txt", "examples/structure.mp4", "examples/first.png"]],
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,