Spaces:
Sleeping
Sleeping
| """Movetoy: a tiny from-scratch text-to-video model, running on a ZeroGPU Space (CPU fallback for local use).""" | |
| import os | |
| import tempfile | |
| import gradio as gr | |
| import imageio | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| try: | |
| import spaces # ZeroGPU: a GPU is attached only while a decorated call runs | |
| gpu = spaces.GPU(duration=lambda prompt, seed, steps, guidance, seconds, progress=None: 10 + 4 * int(seconds)) | |
| except ImportError: # running locally without the `spaces` package | |
| def gpu(fn): | |
| return fn | |
| from tiny_video.model import load_checkpoint, null_tokens, to_uint8, tokenizer_for | |
| torch.set_num_threads(max(1, os.cpu_count() or 1)) | |
| MODEL_REPO = os.environ.get("MOVETOY_REPO", "Kompaktion/Movetoy") | |
| DEVICE = "cuda" if (torch.cuda.is_available() or "spaces" in globals()) else "cpu" | |
| model, _ = load_checkpoint(hf_hub_download(MODEL_REPO, "movetoy.pt"), "cpu") | |
| model = model.to(DEVICE).eval() | |
| tokenizer = tokenizer_for(model.cfg) | |
| SCALE = 8 # 32x32 native -> 256x256 display, nearest-neighbour (no interpolation) | |
| FPS = 8 | |
| WINDOW = model.cfg.frames # 16 frames = 2 s: everything the model was trained to produce in one go | |
| OVERLAP = 4 # frames carried from one window into the next when stitching longer videos | |
| def sample_window(cond, steps, guidance, gen, known=None, shift=2.0): | |
| """Movetoy's Euler sampler for one 16-frame window. If `known` [3,k,H,W] is given, the first k frames are | |
| pinned to it at every step (re-noised to the current time), so the window continues the previous one.""" | |
| cfg = model.cfg | |
| dev = next(model.parameters()).device | |
| x = torch.randn(1, 3, WINDOW, cfg.size, cfg.size, generator=gen).to(dev) # CPU generator: same video on any device | |
| if known is not None: | |
| k = known.shape[1] | |
| eps_k = torch.randn(1, 3, k, cfg.size, cfg.size, generator=gen).to(dev) | |
| u = torch.linspace(1, 0, steps + 1, device=dev) | |
| ts = shift * u / (1 + (shift - 1) * u) | |
| for i in range(steps): | |
| if known is not None: | |
| x[:, :, :k] = (1 - ts[i]) * known + ts[i] * eps_k | |
| with torch.autocast(dev.type, dtype=torch.bfloat16, enabled=dev.type == "cuda"): | |
| v = model(torch.cat([x, x]), ts[i].expand(2), cond).float() | |
| vc, vu = v.chunk(2) | |
| x = x + (ts[i + 1] - ts[i]) * (vu + guidance * (vc - vu)) | |
| if known is not None: | |
| x[:, :, :k] = known | |
| return x.clamp(-1, 1)[0] # [3, 16, H, W] | |
| def generate(prompt, seed, steps, guidance, seconds, progress=gr.Progress()): | |
| """Text -> video. Up to 2 s is one native pass; longer videos stitch 16-frame windows (12 new frames each).""" | |
| prompt = (prompt or "").strip() or "a red circle and blue circle merge into a purple circle" | |
| seed = int(seed) if seed is not None and int(seed) >= 0 else int(np.random.randint(0, 2**31 - 1)) | |
| gen = torch.Generator().manual_seed(seed) | |
| tokens = tokenizer([prompt]) | |
| cond = torch.cat([tokens, null_tokens(tokens, model.cfg.text)]).to(next(model.parameters()).device) | |
| target = int(min(15, max(2, float(seconds))) * FPS) | |
| n_windows = 1 + max(0, -(-(target - WINDOW) // (WINDOW - OVERLAP))) | |
| clip = sample_window(cond, int(steps), float(guidance), gen) | |
| for w in range(1, n_windows): | |
| progress(w / n_windows, desc=f"window {w + 1}/{n_windows}") | |
| nxt = sample_window(cond, int(steps), float(guidance), gen, known=clip[:, -OVERLAP:]) | |
| clip = torch.cat([clip, nxt[:, OVERLAP:]], dim=1) | |
| frames = to_uint8(clip[:, :target].unsqueeze(0))[0].numpy() # [T, 32, 32, 3] | |
| frames = frames.repeat(SCALE, axis=1).repeat(SCALE, axis=2) | |
| path = os.path.join(tempfile.mkdtemp(), "movetoy.mp4") | |
| imageio.mimwrite(path, list(frames), fps=FPS, codec="libx264", quality=9, macro_block_size=1) | |
| note = f"seed {seed} · {len(frames)} frames ({len(frames) / FPS:.1f} s)" | |
| if n_windows > 1: | |
| note += f" · stitched from {n_windows} windows of 16 frames" | |
| return path, note | |
| EXAMPLES = [ | |
| ["a red circle and blue circle merge into a purple circle", 424243, 32, 3.0, 2], | |
| ["a yellow circle and red circle merge into an orange circle", 7, 32, 3.0, 2], | |
| ["a big green star falls down at sunset", 11, 32, 3.0, 2], | |
| ["a pink heart spins in the night sky", 5, 32, 3.0, 2], | |
| ["a small blue square grows bigger on a white background", 3, 32, 3.0, 2], | |
| ["a purple dog dances on the moon", 1, 32, 3.0, 2], | |
| ] | |
| with gr.Blocks(title="Movetoy") as demo: | |
| gr.Markdown( | |
| "# Movetoy\n" | |
| "A **51M-parameter text-to-video model trained from scratch** (random weights, no pretrained parts, " | |
| "no downloaded data) on one RTX 4070. It makes 2-second, 16-frame **32×32** clips of coloured shapes, " | |
| "shown here 8× bigger. It reads *any* text (bytes + hashed words, no vocabulary), but it only *understands* " | |
| "its toy world: colours, shapes, motions, falls, spins, and two-object merges with colour mixing. " | |
| "Words it never saw get you a vibe-based blob.\n\n" | |
| "Runs on ZeroGPU: a 2-second clip takes about a second of GPU. Longer clips (up to 15 s) are stitched from " | |
| "2-second windows that the model continues one after another (it was only ever trained on 2 s), " | |
| "so they drift in funny ways.\n\n" | |
| "**How much training did this take?** ~382,000 procedurally generated captioned clips " | |
| "(23,874 steps × batch 16; every clip unique, rendered on the fly) — about **212 hours** of 32×32 video, " | |
| "~112 minutes on one RTX 4070 — just to learn coloured shapes moving and merging. " | |
| "And it still gets the direction wrong half the time." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="Prompt", value="a red circle and blue circle merge into a purple circle", lines=2) | |
| with gr.Row(): | |
| seed = gr.Number(label="Seed (-1 = random)", value=-1, precision=0) | |
| steps = gr.Slider(8, 64, value=32, step=1, label="Sampling steps") | |
| guidance = gr.Slider(1.0, 6.0, value=3.0, step=0.5, label="Guidance") | |
| seconds = gr.Slider(2, 15, value=2, step=1, label="Length in seconds (over 2 s = stitched, slow on CPU)") | |
| go = gr.Button("Generate", variant="primary") | |
| with gr.Column(): | |
| out = gr.Video(label="Movetoy output (32×32, upscaled)", autoplay=True, loop=True) | |
| used = gr.Markdown() | |
| go.click(generate, [prompt, seed, steps, guidance, seconds], [out, used], api_name="generate") | |
| prompt.submit(generate, [prompt, seed, steps, guidance, seconds], [out, used], api_name=False) | |
| gr.Examples(EXAMPLES, [prompt, seed, steps, guidance, seconds], [out, used], fn=generate, cache_examples=False) | |
| gr.Markdown(f"Model: [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO})") | |
| demo.queue(max_size=32).launch() | |