slightly improved default values

app.py

@@ -3,6 +3,7 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import datasets
 import asyncio
 import numpy as np
+import torch
 
 def make_script(shader_code):
     # code copied and fixed(escaping single quotes to double quotes!!!) from https://webglfundamentals.org/webgl/webgl-shadertoy.html
@@ -272,6 +273,7 @@ outro_text ="""
  - [] support FIM task for better model context
  - [~] include some context for prompt (title, comments before a functions) - now works with the first comment inside a function body (has to be first)
  - [] gradio examples
+ - [] use GPU if available, respect memory restrictions.
 
 ### Notes:
  - this is meant as a resource to show code generation for a "creative" task.
@@ -281,6 +283,14 @@ outro_text ="""
  - 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)
 """
 
+new_shadertoy_code = """void mainImage( out vec4 fragColor, in vec2 fragCoord )
+{
+    // touch the slider to load a shader from the dataset or start coding from here.
+    vec2 uv = fragCoord/iResolution.xy;
+    vec3 col = 0.5 + 0.5*cos(iTime+uv.xyx+vec3(0,2,4));
+    fragColor = vec4(col,1.0);
+}"""
+
 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"]])
@@ -321,9 +331,13 @@ def _parse_functions(in_code):
 PIPE = None
 
 def _make_pipeline(model_cp = "Vipitis/santacoder-finetuned-Shadertoys-fine"): #bad default model for testing
+    # if torch.cuda.is_available():
+    #     device = "cuda"
+    # else:
+    #     device = "cpu"
     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 = pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True) #, device=device)
     PIPE = pipe # set the global?
     print(f"loaded model {model_cp} as a pipline")
     return pipe
@@ -342,7 +356,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, pipeline=PIPE): #default pipeline can't be passed as gloabl?
     """
     Replaces the return statement of a function with a generated one.
     Args:
@@ -356,8 +370,13 @@ def alter_return(orig_code, func_idx="0:", pipeline=PIPE): #default pipeline can
         print("no pipeline found, loading default one")
         pipeline = _make_pipeline()
     
-    print(f"{func_idx=}")
-    func_idx = int(func_idx.split(":")[0].strip())
+    if isinstance(func_idx, str):
+        print(f"{func_idx=}")
+        func_idx = int(func_idx.split(":")[0].strip())
+    elif isinstance(func_idx, int):
+        pass
+    else:
+        raise gr.Error(f"func_idx must be int or str, not {type(func_idx)}")
 
     retrns = []
     retrn_start_idx = orig_code.find("return")
@@ -391,7 +410,7 @@ def _line_chr2char(text, line_idx, chr_idx):
     return char_idx
 
 
-def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE):
+def alter_body(old_code, func_id, funcs_list: list, pipeline=PIPE):
     """
     Replaces the body of a function with a generated one.
     Args:
@@ -401,8 +420,13 @@ def alter_body(old_code, func_id: str, funcs_list: list, pipeline=PIPE):
     Returns:
         str: The altered code.
     """
-    print(f"{func_id=}")
-    func_id = int(func_id.split(":")[0].strip()) #undo their string casting?
+    if isinstance(func_id, str):
+        print(f"{func_id=}")
+        func_id = int(func_id.split(":")[0].strip()) #undo their string casting?
+    elif isinstance(func_id, int):
+        pass
+    else:
+        raise gr.Error(f"func_id must be int or str, not {type(func_id)}")
     func_node = funcs_list[func_id]
     print(f"using for generation: {func_node=}")
     
@@ -468,7 +492,7 @@ 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? 
+    func_dropdown = gr.Dropdown(value=["0: edit the Code (or load a shader) to update this dropdown"], label="chose a function to modify") #breaks if I add a string in before that? #TODO: use type="index" to get int - always gives None?
     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")
@@ -477,12 +501,13 @@ with gr.Blocks() as site:
         with gr.Column():
             source_embed = gr.HTML('<iframe width="640" height="360" frameborder="0" src="" allowfullscreen></iframe>', label="How this shader originally renders")
             our_embed = gr.HTML(label="glsl render of the current code")
-        sample_code = gr.Code("// touch the slider to select a shader", label="Current Code (will update changes you generate)", language=None)
+        sample_code = gr.Code(new_shadertoy_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)
     pipe.value=_make_pipeline("Vipitis/santacoder-finetuned-Shadertoys-fine") # set a default like this?
     funcs = gr.State(value=[])
+    # funcs.value.append(list_dropdown(sample_code.value)[0]) #to circumvent the json issue?
     # hist_state = gr.State(Value={})
     # history_table = gr.JSON()