physicsrob commited on
Commit
3a38e1b
·
verified ·
1 Parent(s): ed03dd6

Add minimal end-to-end example

Browse files
Files changed (1) hide show
  1. example.py +68 -0
example.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Minimal runnable example for this checkpoint.
2
+
3
+ This is complete executable code, not pseudocode. It loads an ordinary
4
+ Hugging Face text-generation checkpoint, runs the bundled E1M1 prompt, executes
5
+ the cursor and pixel commands emitted by the model, and writes ``frame.png``.
6
+ It imports no TorchWright or DOOM implementation.
7
+
8
+ The loop below is all of the post-processing. The transformer emits every
9
+ cursor direction, cursor coordinate, palette index, and run width. This program
10
+ only remembers the cursor, looks up RGB in DOOM's static 256-color palette, and
11
+ asks Pillow to paint those pixels. It contains no map data, geometry, visibility
12
+ tests, texture sampling, lighting calculations, or logic for choosing what gets
13
+ drawn or in what order.
14
+
15
+ This short path is intended to make the mechanism easy to inspect. For
16
+ canonical reproduction, download the checkpoint repository and use:
17
+
18
+ python infer.py --model . --prompt examples/e1m1_prompt.txt --output out
19
+ python tools/txt_to_png.py --input out/output.txt --output out/frame.png
20
+
21
+ That path validates the bundle, preserves the exact generated token IDs,
22
+ records termination and memory information, and validates the decoder inputs.
23
+ See ``README.md`` for details.
24
+ """
25
+
26
+ import json
27
+ from pathlib import Path
28
+
29
+ from huggingface_hub import hf_hub_download
30
+ from PIL import Image
31
+ from transformers import pipeline
32
+
33
+ MODEL = "physicsrob/torchwright-doom-e1m1"
34
+ SCREEN = (320, 200)
35
+
36
+ prompt = Path(hf_hub_download(MODEL, "examples/e1m1_prompt.txt")).read_text()
37
+ colors = json.loads(Path(hf_hub_download(MODEL, "doom_palette.json")).read_text())[
38
+ "colors"
39
+ ]
40
+
41
+ generate = pipeline(
42
+ "text-generation", model=MODEL, device_map="auto", trust_remote_code=False
43
+ )
44
+ output = generate(prompt, return_full_text=False)[0]["generated_text"]
45
+
46
+ image = Image.new("RGB", SCREEN)
47
+ x = y = 0
48
+ advance_x = False
49
+
50
+ for token in output.split():
51
+ command, _, arguments = token.rstrip(")").partition("(")
52
+ if command == "setCursorDirectionX":
53
+ advance_x = True
54
+ elif command == "setCursorDirectionY":
55
+ advance_x = False
56
+ elif command == "setCursorX":
57
+ x = int(arguments)
58
+ elif command == "setCursorY":
59
+ y = int(arguments)
60
+ elif command == "pixel":
61
+ color, width = map(int, arguments.split(","))
62
+ image.paste(tuple(colors[color]), (x, y, x + width, y + 1))
63
+ if advance_x:
64
+ x += width
65
+ else:
66
+ y += 1
67
+
68
+ image.save("frame.png")