Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer | |
| import datasets | |
| import asyncio | |
| import numpy as np | |
| def make_script(shader_code): | |
| # code copied and fixed(escaping single quotes to double quotes!!!) from https://webglfundamentals.org/webgl/webgl-shadertoy.html | |
| script = (""" | |
| <!-- Licensed under a BSD license. See license.html for license --> | |
| <!DOCTYPE html> | |
| <html> | |
| <head> | |
| <meta charset="utf-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=yes"> | |
| <title>WebGL - Shadertoy</title> | |
| <link type="text/css" href="https://webglfundamentals.org/webgl/resources/webgl-tutorials.css" rel="stylesheet" /> | |
| <style> | |
| .divcanvas { | |
| position: relative; | |
| display: inline-block; | |
| } | |
| canvas { | |
| display: block; | |
| } | |
| .playpause { | |
| position: absolute; | |
| left: 10px; | |
| top: 10px; | |
| width: 100%; | |
| height: 100%; | |
| font-size: 60px; | |
| justify-content: center; | |
| align-items: center; | |
| color: rgba(255, 255, 255, 0.3); | |
| transition: opacity 0.2s ease-in-out; | |
| } | |
| .playpausehide, | |
| .playpause:hover { | |
| opacity: 0; | |
| } | |
| .iframe .divcanvas { | |
| display: block; | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="divcanvas"> | |
| <canvas id="canvas"></canvas> | |
| <div class="playpause">▶</div> | |
| </div> | |
| blank canvas here indicates that some of the shadertoy specific functions are not yet supported with this implementation (like #define I believe). you can always copy and paste the code into a shadertoy.com window to try. | |
| </body> | |
| <!-- | |
| for most samples webgl-utils only provides shader compiling/linking and | |
| canvas resizing because why clutter the examples with code thats the same in every sample. | |
| See https://webglfundamentals.org/webgl/lessons/webgl-boilerplate.html | |
| and https://webglfundamentals.org/webgl/lessons/webgl-resizing-the-canvas.html | |
| for webgl-utils, m3, m4, and webgl-lessons-ui. | |
| --> | |
| <script src="https://webglfundamentals.org/webgl/resources/webgl-utils.js"></script> | |
| <script> | |
| "use strict"; | |
| function main() { | |
| // Get A WebGL context | |
| /** @type {HTMLCanvasElement} */ | |
| const canvas = document.querySelector("#canvas"); | |
| const gl = canvas.getContext("webgl"); | |
| if (!gl) { | |
| return; | |
| } | |
| const vs = ` | |
| // an attribute will receive data from a buffer | |
| attribute vec4 a_position; | |
| // all shaders have a main function | |
| void main() { | |
| // gl_Position is a special variable a vertex shader | |
| // is responsible for setting | |
| gl_Position = a_position; | |
| } | |
| `; | |
| const fs = ` | |
| precision highp float; | |
| uniform vec2 iResolution; | |
| uniform vec2 iMouse; | |
| uniform float iTime; | |
| """ + shader_code + """ | |
| void main() { | |
| mainImage(gl_FragColor, gl_FragCoord.xy); | |
| } | |
| `; | |
| // setup GLSL program | |
| const program = webglUtils.createProgramFromSources(gl, [vs, fs]); | |
| // look up where the vertex data needs to go. | |
| const positionAttributeLocation = gl.getAttribLocation(program, "a_position"); | |
| // look up uniform locations | |
| const resolutionLocation = gl.getUniformLocation(program, "iResolution"); | |
| const mouseLocation = gl.getUniformLocation(program, "iMouse"); | |
| const timeLocation = gl.getUniformLocation(program, "iTime"); | |
| // Create a buffer to put three 2d clip space points in | |
| const positionBuffer = gl.createBuffer(); | |
| // Bind it to ARRAY_BUFFER (think of it as ARRAY_BUFFER = positionBuffer) | |
| gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer); | |
| // fill it with a 2 triangles that cover clipspace | |
| gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([ | |
| -1, -1, // first triangle | |
| 1, -1, | |
| -1, 1, | |
| -1, 1, // second triangle | |
| 1, -1, | |
| 1, 1, | |
| ]), gl.STATIC_DRAW); | |
| const playpauseElem = document.querySelector(".playpause"); | |
| const inputElem = document.querySelector(".divcanvas"); | |
| inputElem.addEventListener("mouseover", requestFrame); | |
| inputElem.addEventListener("mouseout", cancelFrame); | |
| let mouseX = 0; | |
| let mouseY = 0; | |
| function setMousePosition(e) { | |
| const rect = inputElem.getBoundingClientRect(); | |
| mouseX = e.clientX - rect.left; | |
| mouseY = rect.height - (e.clientY - rect.top) - 1; // bottom is 0 in WebGL | |
| } | |
| inputElem.addEventListener("mousemove", setMousePosition); | |
| inputElem.addEventListener("touchstart", (e) => { | |
| e.preventDefault(); | |
| playpauseElem.classList.add("playpausehide"); | |
| requestFrame(); | |
| }, {passive: false}); | |
| inputElem.addEventListener("touchmove", (e) => { | |
| e.preventDefault(); | |
| setMousePosition(e.touches[0]); | |
| }, {passive: false}); | |
| inputElem.addEventListener("touchend", (e) => { | |
| e.preventDefault(); | |
| playpauseElem.classList.remove("playpausehide"); | |
| cancelFrame(); | |
| }, {passive: false}); | |
| let requestId; | |
| function requestFrame() { | |
| if (!requestId) { | |
| requestId = requestAnimationFrame(render); | |
| } | |
| } | |
| function cancelFrame() { | |
| if (requestId) { | |
| cancelAnimationFrame(requestId); | |
| requestId = undefined; | |
| } | |
| } | |
| let then = 0; | |
| let time = 0; | |
| function render(now) { | |
| requestId = undefined; | |
| now *= 0.001; // convert to seconds | |
| const elapsedTime = Math.min(now - then, 0.1); | |
| time += elapsedTime; | |
| then = now; | |
| webglUtils.resizeCanvasToDisplaySize(gl.canvas); | |
| // Tell WebGL how to convert from clip space to pixels | |
| gl.viewport(0, 0, gl.canvas.width, gl.canvas.height); | |
| // Tell it to use our program (pair of shaders) | |
| gl.useProgram(program); | |
| // Turn on the attribute | |
| gl.enableVertexAttribArray(positionAttributeLocation); | |
| // Bind the position buffer. | |
| gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer); | |
| // Tell the attribute how to get data out of positionBuffer (ARRAY_BUFFER) | |
| gl.vertexAttribPointer( | |
| positionAttributeLocation, | |
| 2, // 2 components per iteration | |
| gl.FLOAT, // the data is 32bit floats | |
| false, // dont normalize the data | |
| 0, // 0 = move forward size * sizeof(type) each iteration to get the next position | |
| 0, // start at the beginning of the buffer | |
| ); | |
| gl.uniform2f(resolutionLocation, gl.canvas.width, gl.canvas.height); | |
| gl.uniform2f(mouseLocation, mouseX, mouseY); | |
| gl.uniform1f(timeLocation, time); | |
| gl.drawArrays( | |
| gl.TRIANGLES, | |
| 0, // offset | |
| 6, // num vertices to process | |
| ); | |
| requestFrame(); | |
| } | |
| requestFrame(); | |
| requestAnimationFrame(cancelFrame); | |
| } | |
| main(); | |
| </script> | |
| </html> | |
| """) | |
| return script | |
| def make_iframe(shader_code): #keep a single function? | |
| script = make_script(shader_code) | |
| return f"""<iframe width="640" height="512" srcdoc=\'{script}\' allowfullscreen></iframe>""" | |
| intro_text = """ | |
| # Welcome to the interactive shadercoding demo. | |
| This gives you access to a filtered version of the [Shadertoys](https://huggingface.co/datasets/Vipitis/Shadertoys) dataset, only shaders that consist of a single pass are available. | |
| And then lets you use code generation models to make alterations to part of the shadercode. | |
| ## How To Use: | |
| 1. Load any Model for [`text-generation`](https://huggingface.co/models?pipeline_tag=text-generation) and hit ENTER. | |
| 2. Use the slider to sample a shader from the dataset. | |
| - The original shader will be embedding on the left, click on title to get to the source. | |
| - The shadercode will be displayed on the right, this is interactive. | |
| - A preview of the currently displayed shadercode will be displayed on the lower left. (hover to advance time) | |
| 3. use the dropdown to select a function to modify. | |
| 4. press either button to make modifications to that function | |
| 5. you can also edit the code manually. | |
| """ | |
| outro_text =""" | |
| ## Models to try (look at [ShaderEval](https://huggingface.co/spaces/Vipitis/ShaderEval) for an indication of how helpful they will be): | |
| - [gpt2](https://huggingface.co/gpt2) baseline for language models, really struggles with shadercode. | |
| - [bigscience/bloom-1b1](https://huggingface.co/bigscience/bloom-1b1) a newer and larger freely available model. Does understand a big of code. | |
| - [codeparrot/codeparrot-small](https://huggingface.co/codeparrot/codeparrot-small) a model trained on code, but not on shadercode. Manages to graps the patterns. | |
| - [salesforce/codegen-2B-multi](https://huggingface.co/salesforce/codegen-2B-multi) a larger model that indicates some potential. | |
| - [bigcode/santacoder](https://huggingface.co/bigcode/santacoder) a model trained on subset of [TheStack](https://huggingface.co/datasets/bigcode/the-stack), struggles with shadercode. | |
| - [Vipitis/santacoder-finetuned-the-stack-glsl](https://huggingface.co/Vipitis/santacoder-finetuned-the-stack-glsl) fine-tuned by me on the glsl subset of [TheStack](https://huggingface.co/datasets/bigcode/the-stack), is an improvement. | |
| - [Vipitis/santacoder-finetuned-Shadertoys](https://huggingface.co/Vipitis/santacoder-finetuned-Shadertoys) fine-tuned by me on whole shaders from [Shadertoys](https://huggingface.co/datasets/Vipitis/Shadertoys). Does overfit quite a bit with greedy decoding. | |
| - [Vipitis/santacoder-finetuned-Shadertoys-fine](https://huggingface.co/Vipitis/santacoder-finetuned-Shadertoys-fine) fine-tuned by me just functions from [Shadertoys-fine](https://huggingface.co/datasets/Vipitis/Shadertoys-fine). Memorizes the exact function about half the time. | |
| - [bigcode/starcoder](https://huggingface.co/bigcode/starcoder) a very large model which I haven't tried yet. | |
| - **any other model you want to** | |
| ## TODO (feel free to contribute with a [Pull-Request](https://huggingface.co/Vipitis/santacoder-finetuned-the-stack-glsl/discussions?status=open&type=pull_request)): | |
| - [x] use embedded Shadertoy for reference/attribution (done, but some errors) | |
| - [~] working render implementation on CPU only space (as webgl via webglfundamentals, ccs needs fixing for iframe (or hijack Shadertoy iframe)) | |
| - [~] generate variations of return statements [ShaderEval task1](https://huggingface.co/spaces/Vipitis/ShaderEval) (needs to be reworked using the other parts) | |
| - [x] generate whole functions (seems to work quite well) | |
| - [] dropdown for model selection (from curated list or all supported models?) | |
| - [] generation history stating which function and orig/generated returns. (use State ??). do it as comments in the code? | |
| - [] display errros/issues to the user (raise gr.Error could be one idea, but highlighting in the code would be awesome) | |
| - [] generate whole shaders (via prompts guidance, recursive from errors) | |
| - [] accordion with generation parameters (as pipeline_kwargs?) look up starcoder playround and take "inspiration" from there | |
| - [] support FIM task for better model context | |
| - [] gradio examples | |
| ### Notes: | |
| - this is meant as a resource to show code generation for a "creative" task. | |
| - the goal is not to not replace shader artists, but aims to be an assistant instead. | |
| - the space still lacks quite a lot of features, but will continue to evolve. | |
| - this demo can be useful to sannity check evaluation results, where the academic numbers are made. | |
| - If you create a remix with these tools, please attribute the original creator of your starting point when sharing the results. (And perhaps share in the [discussion tab](https://huggingface.co/Vipitis/santacoder-finetuned-the-stack-glsl/discussions?status=open&type=discussion) too) | |
| """ | |
| passes_dataset = datasets.load_dataset("Vipitis/Shadertoys") | |
| single_passes = passes_dataset.filter(lambda x: not x["has_inputs"] and x["num_passes"] == 1) #could also include shaders with no extra functions. | |
| all_single_passes = datasets.concatenate_datasets([single_passes["train"], single_passes["test"]]) | |
| num_samples = len(all_single_passes) | |
| import tree_sitter | |
| from tree_sitter import Language, Parser | |
| Language.build_library("./build/my-languages.so", ['tree-sitter-glsl']) | |
| GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl') | |
| parser = Parser() | |
| parser.set_language(GLSL_LANGUAGE) | |
| def grab_sample(sample_idx): | |
| sample_pass = all_single_passes[sample_idx] | |
| sample_code = sample_pass["code"] | |
| sample_source = sample_pass["source"] | |
| sample_title = sample_pass["title"] | |
| sample_auhtor = sample_pass["author"] | |
| source_iframe = construct_embed(sample_source) | |
| print(f"{source_iframe=}") | |
| # sample_funcs = _parse_functions(sample_code) | |
| # funcs = _parse_functions(sample_code) | |
| # func_identifiers = [f"{idx:2d}: {n.child_by_field_name('declarator').text.decode()}" for idx, n in enumerate(funcs)] | |
| # print(f"updating drop down to:{func_identifiers}") | |
| return sample_pass, sample_code, source_iframe, funcs#, gr.Dropdown.update(choices=func_identifiers) #, sample_title, sample_auhtor | |
| def _parse_functions(in_code): | |
| """ | |
| returns all functions in the code as their actual nodes. | |
| """ | |
| tree = parser.parse(bytes(in_code, "utf8")) | |
| funcs = [n for n in tree.root_node.children if n.type == "function_definition"] | |
| return funcs | |
| PIPE = None | |
| def _make_pipeline(model_cp = "Vipitis/santacoder-finetuned-Shadertoys-fine"): #bad default model for testing | |
| tokenizer = AutoTokenizer.from_pretrained(model_cp, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained(model_cp, trust_remote_code=True) | |
| pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True) | |
| PIPE = pipe # set the global? | |
| print(f"loaded model {model_cp} as a pipline") | |
| return pipe | |
| def process_retn(retn): | |
| return retn.split(";")[0].strip() | |
| def get_full_replacement(orig_code, retn_start_idx, retn_end_idx, prediction) -> str: | |
| """ | |
| Batches the generated return statement into the code and returns the full altered code. | |
| """ | |
| print(f"{orig_code[retn_start_idx:retn_end_idx]=}") | |
| generated = process_retn(prediction) | |
| print(f"{generated=}") | |
| variation = orig_code[:retn_start_idx] + generated + orig_code[retn_end_idx:] | |
| return variation | |
| def alter_return(orig_code, func_idx="0:", pipeline=PIPE): #default pipeline can't be passed as gloabl? | |
| """ | |
| Replaces the return statement of a function with a generated one. | |
| Args: | |
| orig_code (str): The original code. | |
| func_idx (int): The index of the function to replace the return statement of. | |
| pipeline (Pipeline): The pipeline to use for generation. | |
| Returns: | |
| str: The altered code. | |
| """ | |
| if pipeline is None: | |
| print("no pipeline found, loading default one") | |
| pipeline = _make_pipeline() | |
| print(f"{func_idx=}") | |
| func_idx = int(func_idx.split(":")[0].strip()) | |
| retrns = [] | |
| retrn_start_idx = orig_code.find("return") | |
| while retrn_start_idx != -1: | |
| retrn_end_idx = orig_code.find(";", retrn_start_idx) | |
| retrns.append((retrn_start_idx, retrn_end_idx)) | |
| retrn_start_idx = orig_code.find("return", retrn_end_idx) | |
| num_returns = len(retrns) | |
| if num_returns == 0: | |
| print("no return statement found, returning original code") | |
| return orig_code | |
| func_idx = int(max(0, min(func_idx, num_returns - 1))) #clamp to valid range, cast to int as a bodge. | |
| retrn_start_idx, retrn_end_idx = retrns[func_idx] | |
| model_context = orig_code[:retrn_start_idx] #TODO: maximal context? | |
| model_inp = model_context + "return" | |
| new_toks = (retrn_end_idx - retrn_start_idx) * 2 #TODO: approximation, we do have early stopping? maybe also use a number instead? | |
| pipe_generation = pipeline(model_inp, max_new_tokens=new_toks, return_full_text=False)[0]["generated_text"] #pipeline kwargs are missing?! | |
| altered_code = get_full_replacement(orig_code, retrn_start_idx+7, retrn_end_idx, pipe_generation) | |
| return altered_code | |
| def _line_chr2char(text, line_idx, chr_idx): | |
| """ | |
| returns the character index at the given line and character index. | |
| """ | |
| lines = text.split("\n") | |
| char_idx = 0 | |
| for i in range(line_idx): | |
| char_idx += len(lines[i]) + 1 | |
| char_idx += chr_idx | |
| return char_idx | |
| def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE): | |
| """ | |
| Replaces the body of a function with a generated one. | |
| Args: | |
| old_code (str): The original code. | |
| func_node (Node): The node of the function to replace the body of. | |
| pipeline (Pipeline): The pipeline to use for generation. | |
| Returns: | |
| str: The altered code. | |
| """ | |
| print(f"{func_id=}") | |
| func_id = int(func_id.split(":")[0].strip()) #undo their string casting? | |
| func_node = funcs_list[func_id] | |
| print(f"using for generation: {func_node=}") | |
| if pipeline is None: | |
| print("no pipeline found, loading default one") | |
| pipeline = _make_pipeline("Vipitis/santacoder-finetuned-Shadertoys-fine") | |
| func_start_idx = _line_chr2char(old_code, func_node.start_point[0], func_node.start_point[1]) | |
| identifier_str = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() | |
| body_node = func_node.child_by_field_name("body") | |
| body_start_idx = _line_chr2char(old_code, body_node.start_point[0], body_node.start_point[1]) | |
| body_end_idx = _line_chr2char(old_code, body_node.end_point[0], body_node.end_point[1]) | |
| print(f"{old_code[body_start_idx:body_end_idx]=}") | |
| model_context = identifier_str # just this | |
| num_new_tokens = (body_end_idx - body_start_idx) + 10 #TODO: approximation, we do have early stopping? maybe also use a number instead? | |
| print(f"generating up to {num_new_tokens} after {model_context!r}") | |
| generation = pipeline(model_context, max_new_tokens=num_new_tokens, return_full_text=False)[0]["generated_text"] | |
| print(f"{generation=}") | |
| id_with_generation = identifier_str + generation | |
| print(f"{id_with_generation=}") | |
| first_gened_func = _parse_functions(id_with_generation)[0] # truncate generation to a single function? | |
| # strip just the body. | |
| print(f"{first_gened_func=}") | |
| generated_body = first_gened_func.child_by_field_name("body").text.decode() | |
| print(f"{generated_body=}") | |
| altered_code = old_code[:body_start_idx] + generated_body + old_code[body_end_idx:] | |
| return altered_code, pipeline | |
| def add_history(func_id, orig_rtn, gened_rtn, history): | |
| # is this a list? or a JSON dict? | |
| history[func_id] = (orig_rtn, gened_rtn) | |
| return history, history | |
| def list_dropdown(in_code): #only used for auto update, not on sample pick? | |
| funcs = _parse_functions(in_code) | |
| # print(f"updating drop down to:{func_identifiers=}") | |
| func_identifiers = [f"{idx:2d}: {n.child_by_field_name('declarator').text.decode()}" for idx, n in enumerate(funcs)] | |
| # funcs = [n for n in funcs] #wrapped as set to avoid json issues? | |
| print(f"updating drop down to:{func_identifiers}") | |
| return funcs, gr.Dropdown.update(choices=func_identifiers) | |
| def construct_embed(source_url): | |
| shader_id = source_url.split("/")[-1] | |
| return f'<iframe width="640" height="360" frameborder="0" src="https://www.shadertoy.com/embed/{shader_id}?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>' | |
| with gr.Blocks() as site: | |
| top_md = gr.Markdown(intro_text) | |
| model_cp = gr.Textbox(value="Vipitis/santacoder-finetuned-Shadertoys-fine", label="Model Checkpoint (Enter to load!)", interactive=True) | |
| sample_idx = gr.Slider(minimum=0, maximum=num_samples, value=3211, label="pick sample from dataset", step=1.0) | |
| func_dropdown = gr.Dropdown(label="chose a function to modify") #breaks if I add a string in before that? | |
| with gr.Row(): | |
| gen_return_button = gr.Button("generate a alternate return statement", label="generate return") | |
| gen_func_button = gr.Button("generate an alternate function body", label="generate function") | |
| # update_funcs_button = gr.Button("update functions", label="update functions") | |
| with gr.Row(): | |
| with gr.Column(): | |
| source_embed = gr.HTML('<iframe width="640" height="360" frameborder="0" src="https://www.shadertoy.com/embed/WsBcWV?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>', label="How this shader originally renders") | |
| our_embed = gr.HTML(label="glsl render of the current code") | |
| sample_code = gr.Code(label="Current Code (will update changes you generate)", language=None) | |
| bot_md = gr.Markdown(outro_text) | |
| sample_pass = gr.State(value={}) | |
| pipe = gr.State(value=PIPE) | |
| funcs = gr.State(value=[]) | |
| # hist_state = gr.State(Value={}) | |
| # history_table = gr.JSON() | |
| model_cp.submit(fn=_make_pipeline, inputs=[model_cp], outputs=[pipe]) # how can we trigger this on load? | |
| sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_pass, sample_code, source_embed]) | |
| gen_return_button.click(fn=alter_return, inputs=[sample_code, func_dropdown, pipe], outputs=[sample_code]) | |
| gen_func_button.click(fn=alter_body, inputs=[sample_code, func_dropdown, funcs, pipe], outputs=[sample_code, pipe]) | |
| sample_code.change(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]) # to update this after generation, so spans aren't messed up | |
| sample_code.change(fn=make_iframe, inputs=[sample_code], outputs=[our_embed]) #twice could cause issues, find better ways. | |
| site.launch() | |