Update app.py

app.py

@@ -25,20 +25,21 @@ def get_oauth_info(profile: gr.OAuthProfile | None, oauth_token: gr.OAuthToken |
         return print(f'{profile.username}: {org_names}')
 
 
-def compile_model(model_name, sram_size, tensor_size):
+def compile_model(model_name, sram_size, tensor_size, optimize, model_loc):
 
     if oauth_info['token'] is None:
         return "ERROR - please log into HuggingFace to continue"
 
     # Run the comparison
     out_dir = './tmp'
-    model_loc = 'sram'
 
     # Run the model fitter
     results = sr100_model_compiler.sr100_model_compiler(
         model_file=model_name,
         output_dir=f"{out_dir}",
-        model_loc=f"{model_loc}",
+        model_loc=model_loc,
+        optimize=optimize,
+        arena_cache_size=int(float(tensor_size)*1.0e6)
     )
     print(results)
 
@@ -60,27 +61,30 @@ def compile_model(model_name, sram_size, tensor_size):
 
 #def process_data(slider_value):
 #    return slider_value * 2
-
+def update_sliders(sram_slider_value):
+    return gr.update(maximum=sram_slider_value, value=sram_slider_value>>1)
 
 with gr.Blocks() as demo:
     gr.LoginButton()
     text1 = gr.Markdown("SR100 Model Compiler - Compile a tflite model to SR100")
-    
+    user_text = gr.Markdown("")
 
-    # Deploy the two slides
-    sram_slider = gr.Slider(minimum=0, maximum=3, step=0.1, label="Set total SRAM size available in MB", value=3)
-    tensor_slider = gr.Slider(minimum=0, maximum=3, step=0.1, label="Set the SRAM size for tensor calculations in MB", value=1.5)
-    #output_text = gr.Textbox(label="Output")
+    # Setup model inputs
+    with gr.Row():
+        sram_slider = gr.Slider(minimum=0, maximum=3, step=0.1, label="Set total SRAM size available in MB", value=3)
+        tensor_slider = gr.Slider(minimum=0, maximum=3, step=0.1, label="Set the SRAM size for tensor calculations in MB", value=1.5)
+    optimize = gr.Radio(choices=["Performance", "Size"], value='Performance', label='Performance model')
+    model_loc = gr.Radio(choices=["sram", "flash"], value="sram", label='Model weights target')
     
-    model_text = gr.Textbox(label='TFlite model', value='hello_world.tflite')
 
+    # Setup model compile
+    model_text = gr.Textbox(label='TFLite model', value='hello_world.tflite')
     compile_btn = gr.Button("Compile Model")
-    compile_text = gr.Markdown("Waiting for model")
-    user_text = gr.Markdown("")
+    compile_text = gr.Markdown("Waiting for model results")
     
     # Compute options
-    compile_btn.click(compile_model, inputs=[model_text, sram_slider, tensor_slider], outputs=[compile_text])
-    #my_slider.change(fn=process_data, inputs=my_slider, outputs=output_text)
+    compile_btn.click(compile_model, inputs=[model_text, sram_slider, tensor_slider, optimize, model_loc], outputs=[compile_text])
+    sram_slider.change(fn=update_sliders, inputs=sram_slider, outputs=tensor_slider)
     demo.load(get_oauth_info, inputs=None, outputs=user_text)
 
 if __name__ == "__main__":