hardcode mixed precision off for inference

app.py

@@ -150,8 +150,7 @@ def build_interface(estimate_vram_fn):
     return app
         
 
-def estimate_vram(gradio_params):
-    print(gradio_params)
+def estimate_vram(cache_dir, gradio_params):
     model_config = ModelConfig(**filter_params_for_dataclass(ModelConfig, gradio_params))    
     training_config = TrainingConfig(**filter_params_for_dataclass(TrainingConfig, gradio_params))
 
@@ -159,22 +158,17 @@ def estimate_vram(gradio_params):
     # Update model config
     if not gradio_params["repo_id"]:
         return "No model selected!"
-    # If cache directory set, then download config
-    # if gradio_params["cache_dir"]:
-    #     config = scrape_config_from_hub(gradio_params["repo_id"])
-    #     model_config.overwrite_with_hf_config(config)
-    cache_dir="cache/"
+
     # By default, scrape config.json from hub
-    #else:
     config = download_config_from_hub(gradio_params["repo_id"], cache_dir)# gradio_params["cache_dir"])
     model_config.overwrite_with_hf_config(config.to_dict())
 
     if training_config.train:
         total_vram_dict = training_vram_required(model_config, training_config)
-        output_str = f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['gradients']}GB (gradients) + {total_vram_dict['optimizer']}GB (optimizer) + {total_vram_dict['activations']}GB activations"
+        output_str = f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['gradients']}GB (gradients) + {total_vram_dict['optimizer']}GB (optimizer) + {total_vram_dict['activations']}GB (activations)"
     else: # inference
         total_vram_dict = inference_vram_required(model_config, training_config)
-        output_str = f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['kv_cache']}GB (KV cache) + {total_vram_dict['activations']}GB activations"
+        output_str = f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['kv_cache']}GB (KV cache) + {total_vram_dict['activations']}GB (activations)"
     return output_str
 
 if __name__ == "__main__":
@@ -184,7 +178,7 @@ if __name__ == "__main__":
     # Launch gradio interface
     if not args.no_app:
         import gradio as gr
-        estimate_vram_fn = partial(estimate_vram)
+        estimate_vram_fn = partial(estimate_vram, args.cache_dir)
         interface = build_interface(estimate_vram_fn)
         interface.launch()
     # Command line interface

estimate_train_vram.py

@@ -58,6 +58,7 @@ def training_vram_required(model_config, training_config):
 
 
 def inference_vram_required(model_config, training_config):
+    model_config.mixed_precision = False
     # Total inference VRAM = model size + KV cache size + activations + additional overhead
     model_vram = model_memory(parameters=model_config.model_size,
                               precision=model_config.precision,