Update app.py

app.py

@@ -6,7 +6,7 @@ import re
 import json
 import datetime
 from pathlib import Path
-from huggingface_hub import HfApi, hf_hub_download, snapshot_download
+from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file
 import torch
 import subprocess
@@ -14,8 +14,7 @@ import subprocess
 # --- Conversion Function: Safetensors (UNet) → GGUF ---
 def convert_unet_to_gguf(safetensors_path, output_dir, progress=gr.Progress()):
     """
-    Converts a UNet safetensors file to GGUF using gguf-connector's CLI (t2 or t).
-    Assumes the file is named 'unet.safetensors'.
+    Converts a UNet safetensors file to GGUF using gguf-connector's CLI (t2).
     """
     progress(0.1, desc="Starting UNet to GGUF conversion...")
 
@@ -23,59 +22,53 @@ def convert_unet_to_gguf(safetensors_path, output_dir, progress=gr.Progress()):
         # Ensure gguf-connector is available
         import gguf_connector  # noqa
 
-        # Copy input to working dir because ggc expects files in current dir
+        # Prepare working directory
         work_dir = tempfile.mkdtemp()
-        input_path = os.path.join(work_dir, "unet_diffusion_pytorch_model.safetensors")
+        input_name = os.path.basename(safetensors_path)
+        input_path = os.path.join(work_dir, input_name)
         shutil.copy(safetensors_path, input_path)
 
-        # GGUF output will be named automatically like unet.safetensors -> unet.gguf
-        gguf_output_path = os.path.join(work_dir, "unet.gguf")
+        # Expected GGUF output name: same basename + .gguf
+        base_name = os.path.splitext(input_name)[0]
+        gguf_output_path = os.path.join(work_dir, f"{base_name}.gguf")
 
         progress(0.3, desc="Running gguf-connector (t2: safetensors → GGUF)...")
 
-        # Use 'ggc t2' for conversion (beta: unlimited)
-        # This is interactive, so we must simulate input via echo or expect
-        # But since ggc t2 may be interactive, we try non-interactive fallback:
-        # Unfortunately, ggc does not support non-interactive mode robustly.
-        # So we simulate by running in dir and hoping it picks the only file.
-
-        # Change working dir so ggc sees the file
+        # Run ggc t2 in the work directory
         original_cwd = os.getcwd()
         os.chdir(work_dir)
 
         try:
-            # Launch ggc t2 and auto-select first file via input redirection
-            # This is fragile but best-effort
             result = subprocess.run(
                 ["ggc", "t2"],
-                input="1\n",  # select first model
+                input="1\n",  # auto-select first file
                 text=True,
                 capture_output=True,
-                timeout=300
+                timeout=600
             )
             if result.returncode != 0:
                 raise RuntimeError(f"ggc t2 failed: {result.stderr}")
         finally:
             os.chdir(original_cwd)
 
+        # Check if GGUF was created
         if not os.path.exists(gguf_output_path):
-            # Try alternative naming
             candidates = [f for f in os.listdir(work_dir) if f.endswith(".gguf")]
             if not candidates:
                 raise FileNotFoundError("No GGUF file generated by ggc t2")
             gguf_output_path = os.path.join(work_dir, candidates[0])
 
-        # Move to output dir
-        final_gguf_path = os.path.join(output_dir, "unet.gguf")
+        # Move to final output
+        final_gguf_path = os.path.join(output_dir, os.path.basename(gguf_output_path))
         shutil.move(gguf_output_path, final_gguf_path)
 
-        # Also save minimal config
+        # Save minimal config
         config_path = os.path.join(output_dir, "config.json")
         with open(config_path, "w") as f:
             json.dump({
                 "model_type": "unet",
                 "format": "gguf",
-                "source": "converted from safetensors"
+                "source_file": input_name
             }, f)
 
         progress(1.0, desc="Conversion to GGUF complete!")
@@ -88,8 +81,8 @@ def convert_unet_to_gguf(safetensors_path, output_dir, progress=gr.Progress()):
             shutil.rmtree(work_dir, ignore_errors=True)
 
 # --- Main Processing Function ---
-def process_and_upload_unet_to_gguf(repo_url, hf_token, new_repo_id, private_repo, progress=gr.Progress()):
-    if not all([repo_url, hf_token, new_repo_id]):
+def process_and_upload_unet_to_gguf(repo_url, safetensors_filename, hf_token, new_repo_id, private_repo, progress=gr.Progress()):
+    if not all([repo_url, safetensors_filename, hf_token, new_repo_id]):
         return None, "❌ Error: Please fill in all fields.", ""
 
     if not re.match(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$", new_repo_id):
@@ -112,11 +105,11 @@ def process_and_upload_unet_to_gguf(repo_url, hf_token, new_repo_id, private_rep
             return None, "❌ Source must be a Hugging Face model repo.", ""
         src_repo_id = clean_url.replace("https://huggingface.co/", "")
 
-        # Download only unet.safetensors
-        progress(0.15, desc="Downloading unet.safetensors...")
+        # Download specified safetensors file
+        progress(0.15, desc=f"Downloading {safetensors_filename}...")
         safetensors_path = hf_hub_download(
             repo_id=src_repo_id,
-            filename="unet.safetensors",
+            filename=safetensors_filename,
             cache_dir=temp_dir,
             token=hf_token
         )
@@ -146,15 +139,13 @@ tags:
 - converted-by-gradio
 ---
 
-# GGUF UNet Model
+# GGUF Model
 
 Converted from: [`{src_repo_id}`](https://huggingface.co/{src_repo_id})  
-File: `unet.safetensors` → `unet.gguf`
+File: `{safetensors_filename}` → `{os.path.splitext(safetensors_filename)[0]}.gguf`
 
 Converted by: {user_name}  
 Date: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
-
-> ⚠️ This is a GGUF-quantized UNet for storage efficiency. Use with compatible GGUF-aware inference engines.
 """
         with open(os.path.join(output_dir, "README.md"), "w") as f:
             f.write(readme)
@@ -166,13 +157,13 @@ Date: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
             folder_path=output_dir,
             repo_type="model",
             token=hf_token,
-            commit_message="Upload UNet GGUF conversion"
+            commit_message="Upload converted GGUF model"
         )
 
         progress(1.0, desc="✅ Done!")
         result_html = f"""
 ✅ Success!  
-Your GGUF UNet is uploaded to: [{new_repo_id}](https://huggingface.co/{new_repo_id})  
+Your GGUF model is uploaded to: [{new_repo_id}](https://huggingface.co/{new_repo_id})  
 Visibility: {'Private' if private_repo else 'Public'}
 """
         return gr.HTML(result_html), "✅ Conversion and upload successful!", ""
@@ -184,16 +175,21 @@ Visibility: {'Private' if private_repo else 'Public'}
         shutil.rmtree(output_dir, ignore_errors=True)
 
 # --- Gradio UI ---
-with gr.Blocks(title="UNet Safetensors → GGUF Converter") as demo:
-    gr.Markdown("# 🔄 UNet (Safetensors) to GGUF Converter")
-    gr.Markdown("Converts `unet.safetensors` from a Hugging Face model repo to GGUF format for compact storage.")
+with gr.Blocks(title="Safetensors → GGUF Converter") as demo:
+    gr.Markdown("# 🔄 Safetensors to GGUF Converter")
+    gr.Markdown("Converts any `.safetensors` file from a Hugging Face model repo to GGUF format.")
 
     with gr.Row():
         with gr.Column():
             repo_url = gr.Textbox(
                 label="Source Model Repository URL",
                 placeholder="https://huggingface.co/Yabo/FramePainter",
-                info="Must contain 'unet.safetensors'"
+                info="Hugging Face model repo containing your safetensors file"
+            )
+            safetensors_filename = gr.Textbox(
+                label="Safetensors Filename",
+                placeholder="unet.safetensors",
+                info="Name of the .safetensors file in the repo (e.g., unet.safetensors)"
             )
             hf_token = gr.Textbox(
                 label="Hugging Face Token",
@@ -203,7 +199,7 @@ with gr.Blocks(title="UNet Safetensors → GGUF Converter") as demo:
         with gr.Column():
             new_repo_id = gr.Textbox(
                 label="New Repository ID",
-                placeholder="your-username/framepainter-unet-gguf",
+                placeholder="your-username/my-model-gguf",
                 info="Format: username/model-name"
             )
             private_repo = gr.Checkbox(label="Make Private", value=False)
@@ -216,16 +212,16 @@ with gr.Blocks(title="UNet Safetensors → GGUF Converter") as demo:
 
     convert_btn.click(
         fn=process_and_upload_unet_to_gguf,
-        inputs=[repo_url, hf_token, new_repo_id, private_repo],
+        inputs=[repo_url, safetensors_filename, hf_token, new_repo_id, private_repo],
         outputs=[repo_link_output, status_output],
         show_progress=True
     )
 
     gr.Examples(
         examples=[
-            ["https://huggingface.co/Yabo/FramePainter"]
+            ["https://huggingface.co/Yabo/FramePainter", "unet.safetensors"]
         ],
-        inputs=[repo_url]
+        inputs=[repo_url, safetensors_filename]
     )
 
 demo.launch()