Vipitis commited on
Commit
bd74530
1 Parent(s): e925312

adding app and requirements

Browse files
Files changed (2) hide show
  1. app.py +188 -0
  2. requirements.txt +8 -0
app.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
3
+ from wgpu.utils.shadertoy import *
4
+ from wgpu.gui.offscreen import WgpuCanvas as OffscreenCanvas, run as run_offscreen
5
+ import wgpu
6
+ import time
7
+ import ctypes
8
+ import datasets
9
+ from PIL import Image
10
+ import asyncio
11
+ import numpy as np
12
+
13
+ # reimplement the Shadertoy class with offscreen canvas!
14
+ class ShadertoyCustom(Shadertoy):
15
+ def __init__(self, shader_code, resolution=(800, 450), canvas_class=WgpuCanvas, run_fn=run):
16
+ self._canvas_class = canvas_class
17
+ self._fun_fn = run_fn
18
+ super().__init__(shader_code, resolution)
19
+ self._uniform_data = UniformArray(
20
+ ("mouse", "f", 4),
21
+ ("resolution", "f", 3),
22
+ ("time", "f", 1),
23
+ ("time_delta", "f", 1),
24
+ ("frame", "I", 1),
25
+ )
26
+
27
+ self._shader_code = shader_code
28
+ self._uniform_data["resolution"] = resolution + (1,)
29
+
30
+ self._prepare_render()
31
+ self._bind_events()
32
+
33
+ def _prepare_render(self):
34
+ import wgpu.backends.rs # noqa
35
+
36
+ self._canvas = self._canvas_class(title="Shadertoy", size=self.resolution, max_fps=60)
37
+
38
+ adapter = wgpu.request_adapter(
39
+ canvas=self._canvas, power_preference="high-performance"
40
+ )
41
+ self._device = adapter.request_device()
42
+
43
+ self._present_context = self._canvas.get_context()
44
+
45
+ # We use "bgra8unorm" not "bgra8unorm-srgb" here because we want to let the shader fully control the color-space.
46
+ self._present_context.configure(
47
+ device=self._device, format=wgpu.TextureFormat.bgra8unorm
48
+ )
49
+
50
+ shader_type = self.shader_type
51
+ if shader_type == "glsl":
52
+ vertex_shader_code = vertex_code_glsl
53
+ frag_shader_code = (
54
+ builtin_variables_glsl + self.shader_code + fragment_code_glsl
55
+ )
56
+ elif shader_type == "wgsl":
57
+ vertex_shader_code = vertex_code_wgsl
58
+ frag_shader_code = (
59
+ builtin_variables_wgsl + self.shader_code + fragment_code_wgsl
60
+ )
61
+
62
+ vertex_shader_program = self._device.create_shader_module(
63
+ label="triangle_vert", code=vertex_shader_code
64
+ )
65
+ frag_shader_program = self._device.create_shader_module(
66
+ label="triangle_frag", code=frag_shader_code
67
+ )
68
+
69
+ self._uniform_buffer = self._device.create_buffer(
70
+ size=self._uniform_data.nbytes,
71
+ usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST,
72
+ )
73
+
74
+ bind_group_layout = self._device.create_bind_group_layout(
75
+ entries=binding_layout
76
+ )
77
+
78
+ self._bind_group = self._device.create_bind_group(
79
+ layout=bind_group_layout,
80
+ entries=[
81
+ {
82
+ "binding": 0,
83
+ "resource": {
84
+ "buffer": self._uniform_buffer,
85
+ "offset": 0,
86
+ "size": self._uniform_data.nbytes,
87
+ },
88
+ },
89
+ ],
90
+ )
91
+
92
+ self._render_pipeline = self._device.create_render_pipeline(
93
+ layout=self._device.create_pipeline_layout(
94
+ bind_group_layouts=[bind_group_layout]
95
+ ),
96
+ vertex={
97
+ "module": vertex_shader_program,
98
+ "entry_point": "main",
99
+ "buffers": [],
100
+ },
101
+ primitive={
102
+ "topology": wgpu.PrimitiveTopology.triangle_list,
103
+ "front_face": wgpu.FrontFace.ccw,
104
+ "cull_mode": wgpu.CullMode.none,
105
+ },
106
+ depth_stencil=None,
107
+ multisample=None,
108
+ fragment={
109
+ "module": frag_shader_program,
110
+ "entry_point": "main",
111
+ "targets": [
112
+ {
113
+ "format": wgpu.TextureFormat.bgra8unorm,
114
+ "blend": {
115
+ "color": (
116
+ wgpu.BlendFactor.one,
117
+ wgpu.BlendFactor.zero,
118
+ wgpu.BlendOperation.add,
119
+ ),
120
+ "alpha": (
121
+ wgpu.BlendFactor.one,
122
+ wgpu.BlendFactor.zero,
123
+ wgpu.BlendOperation.add,
124
+ ),
125
+ },
126
+ },
127
+ ],
128
+ },
129
+ )
130
+
131
+ def show(self, time: float = 0.0):
132
+ self._canvas.request_draw(self._draw_frame)
133
+ self._fun_fn()
134
+
135
+ text = """
136
+ # Welcome to the interactive shadercoding demo.
137
+ ## (WIP), you can try and explore the dataset a bit right now. (frames are rendered on the fly, not part of the dataset(yet))
138
+
139
+ This gives you access to a filtered version of the [Shadertoys](https://huggingface.co/datasets/Vipitis/Shadertoys) dataset, only shaders that const of a single pass (and have at least one fuction with a return statement) are available.
140
+ In the near future there will be some buttons and sliders to generate variations of the shadercode itself, and hence get some different images.
141
+ If I find an efficient way, the shaders might run in real time and be interactive.
142
+ """
143
+ passes_dataset = datasets.load_dataset("Vipitis/Shadertoys")
144
+ single_passes = passes_dataset.filter(lambda x: not x["has_inputs"] and x["num_passes"] == 1 and x["code"].count("return") >= 1) #filter easier than having a custom loader script?
145
+ all_single_passes = datasets.concatenate_datasets([single_passes["train"], single_passes["test"]])
146
+ num_samples = len(all_single_passes)
147
+
148
+ async def get_image(code, time= 0.0, resolution=(512, 420)):
149
+ shader = ShadertoyCustom(code, resolution, OffscreenCanvas, run_offscreen) #pass offscreen canvas here.
150
+ shader._uniform_data["time"] = time #set any time you want
151
+ shader._canvas.request_draw(shader._draw_frame)
152
+ # frame = shader._canvas.snapshot().data
153
+ frame = np.asarray(shader._canvas.draw())
154
+ img = Image.fromarray(frame)
155
+ # remove transparent pixels
156
+ img = img.convert('RGB')
157
+ return img
158
+
159
+ def grab_sample(sample_idx):
160
+ sample_pass = all_single_passes[sample_idx]
161
+ sample_code = sample_pass["code"]
162
+ sample_source = sample_pass["source"]
163
+ sample_title = sample_pass["title"]
164
+ sample_auhtor = sample_pass["author"]
165
+ return sample_code, sample_source #, sample_title, sample_auhtor
166
+
167
+ def _make_pipeline(model_cp):
168
+ tokenizer = AutoTokenizer.from_pretrained(model_cp, trust_remote_code=True)
169
+ model = AutoModelForCausalLM.from_pretrained(model_cp, trust_remote_code=True)
170
+ return pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True)
171
+
172
+ with gr.Blocks() as site:
173
+ text_md = gr.Markdown(text)
174
+ model_cp = gr.Textbox(value="Vipitis/santacoder-finetuned-Shadertoys", label="Model Checkpoint", interactive=True)
175
+ sample_idx = gr.Slider(minimum=0, maximum=num_samples, value=5, label="pick sample from dataset", step=1.0)
176
+ # run_button = gr.Button(label="generate code")
177
+ render_button = gr.Button("render frame0",label="render frame")
178
+ time_slider = gr.Slider(minimum=0, maximum=10, value=0, label="time (update on release)", step=0.02)
179
+ #output = gr.Textbox(label="Output")
180
+ rendered_frame = gr.Image(shape=(512, 420), label=f"rendered frame preview")
181
+ info_md = gr.Markdown(value="code_source", label="source URL for this shader", interactive=False)
182
+ sample_code = gr.Code(label="Sample Code", language=None, readonly=True, lines=20)
183
+ sample_pass = gr.State(value=None)
184
+ sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_code, info_md])
185
+ time_slider.release(fn=lambda code, time: asyncio.run(get_image(code, time)), inputs=[sample_code, time_slider], outputs=rendered_frame)
186
+ render_button.click(fn=lambda code: asyncio.run(get_image(code)), inputs=[sample_code], outputs=rendered_frame)
187
+ # run_button.click(fn=print, inputs=[model_cp, sample_idx], outputs=output)
188
+ site.launch()
requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ transformers
2
+ datasets
3
+ jsonlines
4
+ wgpu
5
+ torch
6
+ pillow
7
+ gradio
8
+ numpy