Update app.py

app.py

@@ -419,7 +419,7 @@ def cleanup_temp_files(directory=None):
     except Exception as e:
         print(f"Warning: Error during cleanup: {e}")
 
-def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output_widget):
+def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, use_xformers, hf_token, orgs_name, model_name, make_private):
     """Convert the model between different formats."""
     temp_dir = None
     history = ConversionHistory()
@@ -453,6 +453,7 @@ def convert_model(model_to_load, save_precision_as, epoch, global_step, referenc
         args.global_step = global_step
         args.reference_model = reference_model
         args.fp16 = fp16
+        args.use_xformers = use_xformers
 
         update_progress(output_widget, "🔍 Validating input model...", 10)
         args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")
@@ -618,7 +619,7 @@ def get_auto_optimization_suggestions(model_path: str, precision: str, available
     
     return suggestions
 
-def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private, output_widget):
+def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private):
     """Uploads a model to the Hugging Face Hub."""
     try:
         # Login to Hugging Face
@@ -658,15 +659,15 @@ def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_priv
 
 # ---------------------- GRADIO INTERFACE ----------------------
 
-def main(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, hf_token, orgs_name, model_name, make_private):
+def main(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, use_xformers, hf_token, orgs_name, model_name, make_private):
   """Main function orchestrating the entire process."""
   output = gr.Markdown()
 
   # Create tempdir, will only be there for the function
   with tempfile.TemporaryDirectory() as output_path:
-    conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output)
+    conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, use_xformers, output)
 
-    upload_output = upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private, output)
+    upload_output = upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private)
 
     # Return a combined output
     return f"{conversion_output}\n\n{upload_output}"
@@ -750,6 +751,12 @@ with gr.Blocks(css="#main-container { display: flex; flex-direction: column; hei
                 info="Load model in half precision (recommended for CPU usage)"
             )
 
+            use_xformers = gr.Checkbox(
+                label="Enable Memory-Efficient Attention",
+                value=False,
+                info="Enable xFormers for reduced memory usage during conversion"
+            )
+
             # Hugging Face Upload Section
             gr.Markdown("### Upload to Hugging Face (Optional)")
             
@@ -790,6 +797,7 @@ with gr.Blocks(css="#main-container { display: flex; flex-direction: column; hei
                     global_step,
                     reference_model,
                     fp16,
+                    use_xformers,
                     hf_token,
                     orgs_name,
                     model_name,