Update app.py

app.py

@@ -12,6 +12,7 @@ import mlx_vlm
 import mlx.core as mx
 from safetensors import safe_open
 
+# Se utiliza el token de HF para el cache inicial (si se tiene configurado)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 os.environ["HF_HUB_CACHE"] = "cache"
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
@@ -117,36 +118,14 @@ def modify_tensors(mlx_path: str) -> None:
             logging.error(f"Failed to save modified {weights_file}: {e}")
     logging.info("Tensor modification complete.")
 
-def get_user_models(oauth_token: gr.OAuthToken):
-    """
-    Obtiene la lista de modelos del usuario autenticado mediante su token.
-    """
-    if oauth_token is None or oauth_token.token is None:
-        return []
-    try:
-        user_info = whoami(oauth_token.token)
-        username = user_info["name"]
-        api = HfApi(token=oauth_token.token)
-        models = api.list_models(author=username)
-        # Devuelve la lista de model IDs
-        return [model.modelId for model in models]
-    except Exception as e:
-        print(f"Error fetching user models: {e}")
-        return []
-
-def process_model(model_id, q_method, oauth_token: gr.OAuthToken, user_model):
-    """
-    Procesa el modelo a convertir.
-    
-    Si se selecciona un modelo en 'My Models', se usa ese valor; de lo contrario, se utiliza el modelo obtenido desde el buscador.
-    """
+def process_model(model_id, q_method, oauth_token: gr.OAuthToken | None):
     if oauth_token is None or oauth_token.token is None:
         return "You must be logged in to use MLX-my-repo", "error.png"
-    # Priorizar el modelo del usuario si está seleccionado
-    chosen_model = user_model if user_model and user_model != "" else model_id
-    if not chosen_model:
-        return "No model selected.", "error.png"
-    model_name = chosen_model.split("/")[-1]
+    
+    # Se actualiza la variable de entorno HF_TOKEN para que al descargar modelos gated se use la key del perfil
+    os.environ["HF_TOKEN"] = oauth_token.token
+
+    model_name = model_id.split("/")[-1]
     try:
         username = whoami(oauth_token.token)["name"]
     except Exception as e:
@@ -157,29 +136,29 @@ def process_model(model_id, q_method, oauth_token: gr.OAuthToken, user_model):
             if q_method == "FP16":
                 upload_repo = f"{username}/{model_name}-mlx-fp16"
                 try:
-                    mlx_lm.convert(chosen_model, mlx_path=mlx_path, quantize=False, dtype="float16")
+                    mlx_lm.convert(model_id, mlx_path=mlx_path, quantize=False, dtype="float16")
                 except Exception as e1:
                     try:
-                        mlx_vlm.convert(chosen_model, mlx_path=mlx_path, quantize=False, dtype="float16")
+                        mlx_vlm.convert(model_id, mlx_path=mlx_path, quantize=False, dtype="float16")
                     except Exception as e2:
-                        mlx_lm.convert(chosen_model, mlx_path=mlx_path, quantize=False, dtype="float16")
+                        mlx_lm.convert(model_id, mlx_path=mlx_path, quantize=False, dtype="float16")
             else:
                 q_bits = QUANT_PARAMS[q_method]
                 upload_repo = f"{username}/{model_name}-mlx-{q_bits}Bit"
                 try:
-                    mlx_lm.convert(chosen_model, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
+                    mlx_lm.convert(model_id, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
                 except Exception as e1:
                     try:
-                        mlx_vlm.convert(chosen_model, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
+                        mlx_vlm.convert(model_id, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
                     except Exception as e2:
-                        mlx_lm.convert(chosen_model, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
+                        mlx_lm.convert(model_id, mlx_path=mlx_path, quantize=True, q_bits=q_bits)
             try:
                 modify_tensors(mlx_path)
             except FileNotFoundError as e:
                 return f"Error modifying tensors: {e}", "error.png"
             except Exception as e:
                 return f"Error during tensor modification: {e}", "error.png"
-            upload_to_hub(path=mlx_path, upload_repo=upload_repo, hf_path=chosen_model, oauth_token=oauth_token)
+            upload_to_hub(path=mlx_path, upload_repo=upload_repo, hf_path=model_id, oauth_token=oauth_token)
             return (f'Find your repo <a href="https://huggingface.co/{upload_repo}" target="_blank" style="text-decoration:underline">here</a>', "llama.png")
     except Exception as e:
         return f"Error: {e}", "error.png"
@@ -190,33 +169,11 @@ def process_model(model_id, q_method, oauth_token: gr.OAuthToken, user_model):
 css = """.gradio-container { overflow-y: auto; }"""
 with gr.Blocks(css=css) as demo:
     gr.Markdown("You must be logged in to use MLX-my-repo.")
-    # Botón de login que devuelve el token OAuth
-    token_input = gr.LoginButton(min_width=250)
-    
-    # Botón y dropdown para cargar los modelos del usuario
-    load_button = gr.Button("Load My Models")
-    user_models = gr.Dropdown(choices=[], label="My Models", info="List of your Hugging Face models", interactive=True)
-    load_button.click(fn=get_user_models, inputs=token_input, outputs=user_models)
-    
-    # Buscador de modelos en Hugging Face
+    # Se utiliza el LoginButton para obtener el token del usuario, el cual se usará para acceder a modelos gated
+    gr.LoginButton(min_width=250)
     model_id = HuggingfaceHubSearch(label="Hub Model ID", placeholder="Search for model id on Huggingface", search_type="model")
-    q_method = gr.Dropdown(
-        ["FP16", "Q2", "Q3", "Q4", "Q6", "Q8"], 
-        label="Conversion Method", 
-        info="MLX conversion type (FP16 for float16, Q2–Q8 for quantized models)", 
-        value="Q4", filterable=False, visible=True
-    )
-    
-    # La función process_model ahora recibe:
-    # model_id (buscador), q_method, token y user_models (dropdown)
-    iface = gr.Interface(
-        fn=process_model, 
-        inputs=[model_id, q_method, token_input, user_models], 
-        outputs=[gr.Markdown(label="output"), gr.Image(show_label=False)], 
-        title="Create your own MLX Models, blazingly fast ⚡!", 
-        description="The space takes an HF repo as an input, converts it to MLX format (FP16 or quantized), and creates a Public/Private repo under your HF user namespace.", 
-        api_name=False
-    )
+    q_method = gr.Dropdown(["FP16", "Q2", "Q3", "Q4", "Q6", "Q8"], label="Conversion Method", info="MLX conversion type (FP16 for float16, Q2–Q8 for quantized models)", value="Q4", filterable=False, visible=True)
+    iface = gr.Interface(fn=process_model, inputs=[model_id, q_method], outputs=[gr.Markdown(label="output"), gr.Image(show_label=False)], title="Create your own MLX Models, blazingly fast ⚡!", description="The space takes an HF repo as an input, converts it to MLX format (FP16 or quantized), and creates a Public/Private repo under your HF user namespace.", api_name=False)
 
 def restart_space():
     try: