some cleaning

optimum_neuron_export.py

@@ -39,7 +39,6 @@ from optimum.neuron import (
     NeuronModelForSeq2SeqLM,
 )
 
-# Import diffusers pipelines
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
@@ -339,9 +338,9 @@ def export_decoder_model(model_id: str, folder: str, token:str) -> Generator:
         raise
     
 def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str, pipeline_name: str = None) -> Any:
-
-    operations = []
+    """Export model and git add it."""
     
+    operations = []
     try:
         if model_type == "diffusers":
             export_gen = export_diffusion_model(model_id, pipeline_name, task_or_pipeline, folder, token) 
@@ -372,7 +371,6 @@ def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, fo
             readme_path = os.path.join(folder, "README.md")
             card.save(readme_path)
             
-            # Check if README.md is already in operations, if so update, else add
             readme_op = next((op for op in operations if op.path_in_repo == "README.md"), None)
             if readme_op:
                 readme_op.path_or_fileobj = readme_path
@@ -390,6 +388,7 @@ def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, fo
 
 def generate_neuron_repo_name(api, original_model_id: str, task_or_pipeline: str, token:str) -> str:
     """Generate a name for the Neuron-optimized repository."""
+    
     requesting_user = api.whoami(token=token)["name"]
     base_name = original_model_id.replace('/', '-')
     return f"{requesting_user}/{base_name}-neuron"
@@ -403,11 +402,9 @@ def create_neuron_repo_and_upload(
     token: str,
     pipeline_name: str = None,
 ) -> Generator[Union[str, RepoUrl], None, None]:
-    """
-    Creates a new repository with Neuron files and uploads them.
-    """
+    """Creates a new repository with Neuron files and uploads them."""
+    
     api = HfApi(token=token)
-
     if task_or_pipeline == "auto" and model_type == "transformers":
         try:
             task_or_pipeline = TasksManager.infer_task_from_model(original_model_id, token=token)
@@ -418,7 +415,6 @@ def create_neuron_repo_and_upload(
     neuron_repo_name = generate_neuron_repo_name(api, original_model_id, task_or_pipeline, token)
     
     try:
-        # Create the repository
         repo_url = create_repo(
             repo_id=neuron_repo_name,
             token=token,
@@ -427,7 +423,6 @@ def create_neuron_repo_and_upload(
             exist_ok=True,  
         )
                 
-        # Get the appropriate class name for the Python example
         if model_type == "transformers":
             model_class = TASK_TO_MODEL_CLASS.get(task_or_pipeline)
             model_class_name = model_class.__name__ if model_class else "NeuronModel"
@@ -492,12 +487,10 @@ For the original model, training details, and more information, please visit: [{
         # Update the README in operations
         readme_op = next((op for op in operations if op.path_in_repo == "README.md"), None)
         if readme_op:
-            # Create a temporary file with the new content
             with NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
                 f.write(neuron_readme_content)
                 readme_op.path_or_fileobj = f.name
         else:
-            # Add new README operation
             with NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
                 f.write(neuron_readme_content)
                 operations.append(CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=f.name))
@@ -535,11 +528,9 @@ def create_readme_pr_for_original_model(
     requesting_user: str,
     token: str,
 ) -> Generator[Union[str, CommitInfo], None, None]:
-    """
-    Creates a PR on the original model repository to add a link to the Neuron-optimized version.
-    """
+    """Creates a PR on the original model repository to add a link to the Neuron-optimized version."""
+    
     api = HfApi(token=token)
-
     yield f"📝 Creating PR to add Neuron repo link in {original_model_id}..."
 
     try:
@@ -561,7 +552,6 @@ def create_readme_pr_for_original_model(
             with open(current_readme_path, 'r', encoding='utf-8') as f:
                 readme_content = f.read()
         except Exception:
-            # If README doesn't exist, create a basic one
             readme_content = f"# {original_model_id}\n\n"
 
         # Add Neuron optimization section, separated by a horizontal rule
@@ -626,11 +616,9 @@ def upload_to_custom_repo(
     requesting_user: str,
     token: str,
 ) -> Generator[Union[str, CommitInfo], None, None]:
-    """
-    Uploads neuron files to a custom repository and creates a PR.
-    """
-    api = HfApi(token=token)
+    """Uploads neuron files to a custom repository and creates a PR."""
     
+    api = HfApi(token=token)
     try:
         # Ensure the custom repo exists
         api.repo_info(repo_id=custom_repo_id, repo_type="model")
@@ -702,7 +690,7 @@ def convert(
         result_info = {}
 
         try:
-            # --- Export Logic ---
+            # Export Logic 
             export_gen = export_and_git_add(model_id, task_or_pipeline, model_type, export_folder, token=token, pipeline_name=pipeline_name)
             operations = None
             for message in export_gen:
@@ -715,7 +703,7 @@ def convert(
             if not operations:
                 raise Exception("Export process did not produce any files to commit.")
 
-            # --- Cache Handling ---
+            # Cache Handling
             if pr_options.get("create_cache_pr"):
                 yield "0", f"📤 Creating a Pull Request for the cache repository ..."
                 
@@ -731,7 +719,7 @@ This PR adds newly compiled cache artifacts for the model:
 These files contain precompiled Neuron-optimized representations of the model, allowing faster loading and inference on AWS Neuron hardware.
 """ 
                     
-                    # 1. Create an instance of your generator
+                    # Create an instance of your generator
                     commit_message = f"Synchronizing local compiler cache of {model_id}"
                     inputs = get_default_inputs(task_or_pipeline, pipeline_name)
                     commit_description = f"""
@@ -755,7 +743,7 @@ These files contain precompiled Neuron-optimized representations of the model, a
                     )
 
                     pr_url = None
-                    # 2. Loop to process yielded status messages and capture the final return value
+                    # Loop to process yielded status messages and capture the final return value
                     while True:
                         try:
                             # Get the next status message from your generator
@@ -766,7 +754,7 @@ These files contain precompiled Neuron-optimized representations of the model, a
                             pr_url = e.value
                             break # Exit the loop
 
-                    # 3. Process the final result
+                    # Process the final result
                     if pr_url:
                         yield "0", f"✅ Successfully captured PR URL."
                         result_info["cache_pr"] = pr_url
@@ -776,11 +764,10 @@ These files contain precompiled Neuron-optimized representations of the model, a
                 except Exception as e:
                     yield "0", f"❌ Failed to create cache PR: {e}"
             
-            # --- New Repository Creation (Replaces Model PR) ---
+            # New Repository Creation (Replaces Model PR)
             if pr_options.get("create_neuron_repo"):
                 yield "0", "🏗️ Creating new Neuron-optimized repository..."
                 neuron_repo_url = None
-                # Generate the repo name first so we can use it consistently
                 neuron_repo_name = generate_neuron_repo_name(api, model_id, task_or_pipeline, token)
                 
                 repo_creation_gen = create_neuron_repo_and_upload(
@@ -809,7 +796,7 @@ These files contain precompiled Neuron-optimized representations of the model, a
                 if readme_pr:
                     result_info["readme_pr"] = f"https://huggingface.co/{model_id}/discussions/{readme_pr.pr_num}"
             
-            # --- Custom Repository PR ---
+            # Custom Repository PR
             if pr_options.get("create_custom_pr"):
                 custom_repo_id = pr_options["custom_repo_id"]
                 yield "0", f"📤 Creating PR in custom repository: {custom_repo_id}..."
@@ -827,5 +814,4 @@ These files contain precompiled Neuron-optimized representations of the model, a
 
         except Exception as e:
             yield "1", f"❌ Conversion failed with a critical error: {e}"
-            # Re-raise the exception to be caught by the outer try-except in the Gradio app if needed
             raise