remove wgpu, add documentation

app.py

@@ -1,137 +1,9 @@
 import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-from wgpu.utils.shadertoy import * 
-from wgpu.gui.offscreen import WgpuCanvas as OffscreenCanvas, run as run_offscreen
-import wgpu
-import time
-import ctypes
 import datasets
-from PIL import Image
 import asyncio
 import numpy as np
 
-# reimplement the Shadertoy class with offscreen canvas!
-class ShadertoyCustom(Shadertoy):
-    def __init__(self, shader_code, resolution=(800, 450), canvas_class=WgpuCanvas, run_fn=run):
-        self._canvas_class = canvas_class
-        self._fun_fn = run_fn
-        super().__init__(shader_code, resolution)
-        self._uniform_data = UniformArray(
-            ("mouse", "f", 4),
-            ("resolution", "f", 3),
-            ("time", "f", 1),
-            ("time_delta", "f", 1),
-            ("frame", "I", 1),
-        )
-        
-        self._shader_code = shader_code
-        self._uniform_data["resolution"] = resolution + (1,)
-
-        self._prepare_render()
-        self._bind_events()
-    
-    def _prepare_render(self):
-        import wgpu.backends.rs  # noqa
-
-        self._canvas = self._canvas_class(title="Shadertoy", size=self.resolution, max_fps=60)
-
-        adapter = wgpu.request_adapter(
-            canvas=self._canvas, power_preference="high-performance"
-        )
-        self._device = adapter.request_device()
-
-        self._present_context = self._canvas.get_context()
-
-        # We use "bgra8unorm" not "bgra8unorm-srgb" here because we want to let the shader fully control the color-space.
-        self._present_context.configure(
-            device=self._device, format=wgpu.TextureFormat.bgra8unorm
-        )
-
-        shader_type = self.shader_type
-        if shader_type == "glsl":
-            vertex_shader_code = vertex_code_glsl
-            frag_shader_code = (
-                builtin_variables_glsl + self.shader_code + fragment_code_glsl
-            )
-        elif shader_type == "wgsl":
-            vertex_shader_code = vertex_code_wgsl
-            frag_shader_code = (
-                builtin_variables_wgsl + self.shader_code + fragment_code_wgsl
-            )
-
-        vertex_shader_program = self._device.create_shader_module(
-            label="triangle_vert", code=vertex_shader_code
-        )
-        frag_shader_program = self._device.create_shader_module(
-            label="triangle_frag", code=frag_shader_code
-        )
-
-        self._uniform_buffer = self._device.create_buffer(
-            size=self._uniform_data.nbytes,
-            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST,
-        )
-
-        bind_group_layout = self._device.create_bind_group_layout(
-            entries=binding_layout
-        )
-
-        self._bind_group = self._device.create_bind_group(
-            layout=bind_group_layout,
-            entries=[
-                {
-                    "binding": 0,
-                    "resource": {
-                        "buffer": self._uniform_buffer,
-                        "offset": 0,
-                        "size": self._uniform_data.nbytes,
-                    },
-                },
-            ],
-        )
-
-        self._render_pipeline = self._device.create_render_pipeline(
-            layout=self._device.create_pipeline_layout(
-                bind_group_layouts=[bind_group_layout]
-            ),
-            vertex={
-                "module": vertex_shader_program,
-                "entry_point": "main",
-                "buffers": [],
-            },
-            primitive={
-                "topology": wgpu.PrimitiveTopology.triangle_list,
-                "front_face": wgpu.FrontFace.ccw,
-                "cull_mode": wgpu.CullMode.none,
-            },
-            depth_stencil=None,
-            multisample=None,
-            fragment={
-                "module": frag_shader_program,
-                "entry_point": "main",
-                "targets": [
-                    {
-                        "format": wgpu.TextureFormat.bgra8unorm,
-                        "blend": {
-                            "color": (
-                                wgpu.BlendFactor.one,
-                                wgpu.BlendFactor.zero,
-                                wgpu.BlendOperation.add,
-                            ),
-                            "alpha": (
-                                wgpu.BlendFactor.one,
-                                wgpu.BlendFactor.zero,
-                                wgpu.BlendOperation.add,
-                            ),
-                        },
-                    },
-                ],
-            },
-        )
-    
-    def show(self, time: float = 0.0):
-        self._canvas.request_draw(self._draw_frame)
-        self._fun_fn()
-
 def make_script(shader_code):
     # code copied and fixed(escaping single quotes to double quotes!!!) from https://webglfundamentals.org/webgl/webgl-shadertoy.html
     script = ("""
@@ -358,23 +230,54 @@ def make_iframe(shader_code): #keep a single function?
      return f"""<iframe width="640" height="512" srcdoc=\'{script}\' allowfullscreen></iframe>"""
     
 
-text = """
+intro_text = """
 # Welcome to the interactive shadercoding demo.
-## (WIP), you can try and explore the dataset a bit right now. (frames are rendered on the fly, not part of the dataset(yet))
-
-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.
-In the near future there will be some buttons and sliders to generate variations of the shadercode itself, and hence get some different images. 
-If I find an efficient way, the shaders might run in real time and be interactive.
+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.
+"""
 
-## TODO:
+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 (use the browser for WebGPU?, maybe via an iFrame too?) freespace uses lavapipe which works on really simple stuff.
+ - [~] 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)
- - [] generation history stating which function and orig/generated returns. (use State ??). do it as comments in the code?
  - [x] generate whole functions (seems to work quite well)
- - [] display errros/issues to the user
- - [] generate whole shaders (via prompts?)
- - [] accordion with generation parameters (as pipeline_kwargs?)
+ - [] 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")
@@ -389,23 +292,6 @@ GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl')
 parser = Parser()
 parser.set_language(GLSL_LANGUAGE)
 
-
-async def get_image(code, time= 0.0, resolution=(512, 420)):
-    tree = parser.parse(bytes(code, "utf8"))
-    if tree.root_node.has_error:
-        print("ERROR in the tree, aborting.")
-        raise gr.Error("the code seems to have issues")
-        return None
-    shader = ShadertoyCustom(code, resolution, OffscreenCanvas, run_offscreen) #pass offscreen canvas here.
-    shader._uniform_data["time"] = time #set any time you want
-    shader._canvas.request_draw(shader._draw_frame)
-    # frame = shader._canvas.snapshot().data
-    frame = np.asarray(shader._canvas.draw())
-    img = Image.fromarray(frame)
-    # remove transparent pixels 
-    img = img.convert('RGB')
-    return img
-
 def grab_sample(sample_idx):
     sample_pass = all_single_passes[sample_idx]
     sample_code = sample_pass["code"]
@@ -454,7 +340,7 @@ def get_full_replacement(orig_code, retn_start_idx, retn_end_idx, prediction) ->
     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?
+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:
@@ -467,6 +353,9 @@ def alter_return(orig_code, func_idx=0, pipeline=PIPE): #default pipeline can't
     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")
@@ -499,7 +388,8 @@ def _line_chr2char(text, line_idx, chr_idx):
     char_idx += chr_idx
     return char_idx
 
-def alter_body(old_code, func_id, funcs_list, pipeline=PIPE):
+
+def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE):
     """
     Replaces the body of a function with a generated one.
     Args:
@@ -561,7 +451,7 @@ def construct_embed(source_url):
     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:
-    text_md = gr.Markdown(text)
+    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? 
@@ -569,31 +459,24 @@ with gr.Blocks() as site:
         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")
-    render_button = gr.Button("render frame0 (can carsh the sapce on invalid shadercode)",label="render frame")
-    time_slider = gr.Slider(minimum=0, maximum=10, value=0, label="time (update on release, also used to pick other functions as a bodge)", step=0.02)
     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")
-            rendered_frame = gr.Image(shape=(512, 420), label=f"rendered frame preview", type="pil") #colors are messed up?
             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])
-    sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_pass, sample_code, source_embed])
-    # sample_idx.release(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]) #use multiple event handles to call other functions! seems to not work really well. always messes up
-    gen_return_button.click(fn=alter_return, inputs=[sample_code, time_slider, pipe], outputs=[sample_code])
+    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])
-    # run_button.click(fn=add_history, inputs=[time_slider, sample_pass, sample_code, hist_state], outputs=[history_table, hist_state])
-    # sample_idx.release(fn=construct_embed, inputs=[sample_idx], outputs=[source_embed]) #twice to make have different outputs?
     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.
-    time_slider.release(fn=lambda code, time: asyncio.run(get_image(code, time)), inputs=[sample_code, time_slider], outputs=rendered_frame)
-    render_button.click(fn=lambda code: asyncio.run(get_image(code)), inputs=[sample_code], outputs=rendered_frame)
-    # run_button.click(fn=print, inputs=[model_cp, sample_idx], outputs=output)
 site.launch()

requirements.txt

@@ -1,11 +1,8 @@
 transformers
 datasets
 jsonlines
-wgpu
 torch
 pillow
 gradio
-numpy
-glfw
 jupylet
 tree-sitter