Update app.py

app.py

@@ -1,260 +1,141 @@
-import gradio as gr
-from huggingface_hub import HfApi, hf_hub_download, BucketFile, BucketFolder
-from safetensors.torch import load_file, save_file
 import os
-import re
-import time
-import shutil
-import tempfile
-import torch
 import gc
+import torch
+import shutil
+import gradio as gr
+from huggingface_hub import HfApi, hf_hub_download
+from safetensors.torch import load_file, save_file
 
-RATE_LIMIT_PATTERN = re.compile(r"Retry after\s*(\d+)\s*seconds", re.IGNORECASE)
+SOURCE_REPO = "Tongyi-MAI/Z-Image-Turbo"
+TARGET_REPO = "rootlocalghost/Z-Image-Turbo-FP8"
+TEMP_DIR = "temp_processing_dir"
 
-def _is_rate_limit_error(exc: Exception) -> bool:
-    message = str(exc).lower()
-    return "429" in message or "too many requests" in message or "rate limit" in message
+def convert_and_upload(token):
+    if not token:
+        yield "❌ Error: Please provide a valid Hugging Face Write Token."
+        return
 
-def _get_retry_after(exc: Exception, default: int = 2) -> int:
-    match = RATE_LIMIT_PATTERN.search(str(exc))
-    if match:
-        return int(match.group(1))
-    return default
+    api = HfApi(token=token)
+    yield f"🔄 Connecting to Hugging Face and verifying target repo: {TARGET_REPO}..."
 
-def _retry_api_call(fn, *args, retries: int = 3, **kwargs):
-    delay = 2
-    for attempt in range(1, retries + 1):
-        try:
-            return fn(*args, **kwargs)
-        except Exception as exc:
-            if not _is_rate_limit_error(exc) or attempt == retries:
-                raise
-            wait = _get_retry_after(exc, delay)
-            time.sleep(wait)
-            delay = min(delay * 2, 60)
+    # Ensure the target repo exists, create it if it doesn't
+    try:
+        api.create_repo(repo_id=TARGET_REPO, exist_ok=True, private=False)
+    except Exception as e:
+        yield f"❌ Error checking/creating repo: {str(e)}\nMake sure your token has 'Write' permissions."
+        return
 
-def _format_bucket_uri(bucket_id: str) -> str:
-    bucket_id = bucket_id.strip()
-    if bucket_id.startswith("hf://buckets/"):
-        return bucket_id
-    if bucket_id.startswith("buckets/"):
-        return f"hf://{bucket_id}"
-    return f"hf://buckets/{bucket_id}"
+    yield "📋 Fetching file list from the source repository..."
+    try:
+        files = api.list_repo_files(SOURCE_REPO)
+    except Exception as e:
+        yield f"❌ Error fetching files: {str(e)}"
+        return
 
-def _get_target_dtype(precision):
-    if precision == "FP8":
-        return torch.float8_e4m3fn
-    elif precision == "FP16":
-        return torch.float16
-    elif precision == "BF16":
-        return torch.bfloat16
-    return None
+    # Create a temporary directory for safe local processing
+    os.makedirs(TEMP_DIR, exist_ok=True)
 
-def _process_and_upload_file(local_path, file_path, target_repo, repo_type, api, hf_token, precision, target_folders):
-    target_dtype = _get_target_dtype(precision)
-    folders_to_check = [f.strip() for f in target_folders.split(",")] if target_folders else []
-    
-    # Check if this file is in a targeted folder (or if no target folders are specified)
-    in_target_folder = any(f in file_path for f in folders_to_check) if folders_to_check else True
+    for file in files:
+        yield f"⏳ Processing {file}..."
 
-    if precision != "None" and file_path.endswith(".safetensors") and in_target_folder:
-        # Quantize the file
-        tensors = load_file(local_path)
-        for k, v in tensors.items():
-            if v.is_floating_point():
-                tensors[k] = v.to(target_dtype)
-        
-        converted_path = local_path + ".converted"
-        save_file(tensors, converted_path)
-        
-        # Flush memory
-        del tensors
-        gc.collect()
+        try:
+            # Download file locally without using the central symlink cache 
+            # This is critical to prevent the 50GB Space disk from filling up
+            local_path = hf_hub_download(
+                repo_id=SOURCE_REPO,
+                filename=file,
+                local_dir=TEMP_DIR,
+                local_dir_use_symlinks=False 
+            )
+
+            # Check if it's a safetensor file inside the target directories
+            if file.endswith(".safetensors") and ("text_encoder/" in file or "transformer/" in file):
+                yield f"🧠 Quantizing {file} to FP8 (This may take a minute)..."
+                
+                # Load tensors into RAM
+                tensors = load_file(local_path)
 
-        api.upload_file(
-            path_or_fileobj=converted_path,
-            path_in_repo=file_path,
-            repo_id=target_repo,
-            repo_type=repo_type,
-            commit_message=f"clone & quantize ({precision}) {file_path}",
-            token=hf_token,
-        )
-        os.remove(converted_path)
-    else:
-        # Upload as-is
-        api.upload_file(
-            path_or_fileobj=local_path,
-            path_in_repo=file_path,
-            repo_id=target_repo,
-            repo_type=repo_type,
-            commit_message=f"clone {file_path}",
-            token=hf_token,
-        )
+                # Cast all floating point tensors to FP8
+                keys = list(tensors.keys())
+                for k in keys:
+                    if tensors[k].is_floating_point():
+                        tensors[k] = tensors[k].to(torch.float8_e4m3fn)
 
-def _stream_clone_repo(source_repo, target_repo, repo_type, api, hf_token, precision, target_folders):
-    file_paths = api.list_repo_files(
-        repo_id=source_repo,
-        repo_type=repo_type,
-        token=hf_token,
-    )
-    if not file_paths:
-        raise ValueError("source repo is empty or could not be listed")
+                # Save the quantized tensors to a new temp file
+                converted_path = os.path.join(TEMP_DIR, "converted.safetensors")
+                save_file(tensors, converted_path)
+
+                # Wipe the tensors from RAM immediately to stay under the 16GB limit
+                del tensors
+                gc.collect()
 
-    with tempfile.TemporaryDirectory(prefix="hf_file_") as root_dir:
-        for file_path in file_paths:
-            if file_path.endswith("/"):
-                continue
-            try:
-                downloaded_path = hf_hub_download(
-                    repo_id=source_repo,
-                    filename=file_path,
-                    repo_type=repo_type,
-                    local_dir=root_dir,
-                    local_dir_use_symlinks=False,
-                    token=hf_token,
+                yield f"☁️ Uploading FP8 version of {file}..."
+                api.upload_file(
+                    path_or_fileobj=converted_path,
+                    path_in_repo=file,
+                    repo_id=TARGET_REPO,
+                    commit_message=f"Upload FP8 quantized {file}"
                 )
-                if not os.path.isfile(downloaded_path):
-                    raise ValueError(f"Downloaded file not found: {downloaded_path}")
                 
-                _process_and_upload_file(downloaded_path, file_path, target_repo, repo_type, api, hf_token, precision, target_folders)
-
-            finally:
-                if os.path.exists(downloaded_path):
-                    os.remove(downloaded_path)
-            gc.collect()
-
-def _upload_local_source(source_path, target_repo, repo_type, api, hf_token, precision, target_folders):
-    if not os.path.isdir(source_path):
-        raise ValueError("Local source path must be an existing directory.")
-    
-    if precision == "None":
-        # Bulk upload if no quantization is needed
-        api.upload_large_folder(
-            repo_id=target_repo,
-            folder_path=source_path,
-            repo_type=repo_type,
-            num_workers=1,
-            print_report=False,
-        )
-    else:
-        # File-by-file processing for local quantization
-        for root, _, files in os.walk(source_path):
-            for file in files:
-                local_file_path = os.path.join(root, file)
-                repo_file_path = os.path.relpath(local_file_path, source_path).replace("\\", "/")
-                _process_and_upload_file(local_file_path, repo_file_path, target_repo, repo_type, api, hf_token, precision, target_folders)
-                gc.collect()
+                # Clean up the converted file
+                os.remove(converted_path)
+
+            else:
+                yield f"☁️ Copying {file} as-is..."
+                api.upload_file(
+                    path_or_fileobj=local_path,
+                    path_in_repo=file,
+                    repo_id=TARGET_REPO,
+                    commit_message=f"Copy {file} from original repo"
+                )
 
-def _stream_clone_bucket(source_repo, target_repo, repo_type, api, hf_token, precision, target_folders):
-    bucket_uri = _format_bucket_uri(source_repo)
-    bucket_id = bucket_uri[len("hf://"):]
-    items = api.list_bucket_tree(bucket_id=bucket_id, recursive=True, token=hf_token)
-    
-    with tempfile.TemporaryDirectory(prefix="hf_file_") as root_dir:
-        for item in items:
-            if isinstance(item, BucketFolder):
-                continue
-            if isinstance(item, BucketFile):
-                local_path = os.path.join(root_dir, item.path)
-                os.makedirs(os.path.dirname(local_path), exist_ok=True)
-                try:
-                    api.download_bucket_files(
-                        bucket_id=bucket_id,
-                        files=[(item.path, local_path)],
-                        token=hf_token,
-                    )
-                    _process_and_upload_file(local_path, item.path, target_repo, repo_type, api, hf_token, precision, target_folders)
-                finally:
-                    if os.path.exists(local_path):
-                        os.remove(local_path)
+            # Delete the downloaded original file to free up disk space
+            if os.path.exists(local_path):
+                os.remove(local_path)
+            
+            # Final sweep of memory before the next file
             gc.collect()
 
-def stealth_clone_hf_repo(hf_token_ui, source_repo, source_type, target_repo, repo_type, precision, target_folders):
-    # Use UI token first, fallback to environment variable
-    hf_token = hf_token_ui.strip() if hf_token_ui.strip() else os.environ.get("HF_TOKEN")
-    
-    if not hf_token:
-        return "error: HF_TOKEN secret not found and no token provided in the UI."
+        except Exception as e:
+            yield f"⚠️ Error processing {file}: {str(e)}\nSkipping to next file..."
 
-    api = HfApi(token=hf_token)
-    try:
-        _retry_api_call(
-            api.create_repo,
-            repo_id=target_repo,
-            repo_type=repo_type,
-            exist_ok=True,
-        )
-
-        if source_type == "bucket":
-            _stream_clone_bucket(source_repo, target_repo, repo_type, api, hf_token, precision, target_folders)
-        elif source_type == "local":
-            _upload_local_source(source_repo, target_repo, repo_type, api, hf_token, precision, target_folders)
-        else:
-            _stream_clone_repo(source_repo, target_repo, repo_type, api, hf_token, precision, target_folders)
+    # Clean up the processing directory
+    if os.path.exists(TEMP_DIR):
+        shutil.rmtree(TEMP_DIR)
 
-        return f"success! cleanly cloned and processed {source_repo} to {target_repo} with no tags."
-    except Exception as e:
-        return f"error: {type(e).__name__}: {str(e)}"
+    yield "✅ All files processed and successfully uploaded to your repository!"
 
-# Build UI
+# Build the Gradio Web Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## 🥷 Stealth Quantizing Cloner")
+    gr.Markdown("# 🚀 Z-Image-Turbo FP8 Quantizer & Uploader")
     gr.Markdown(
-        "Clone repos, buckets, or local paths cleanly without the 'duplicated from' tag. "
-        "Optionally intercept `.safetensors` files during the transfer and quantize them on the fly."
+        f"This tool sequentially downloads files from `{SOURCE_REPO}`, quantizes the **text_encoder** and **transformer** "
+        f"`.safetensors` files to FP8 (`float8_e4m3fn`), and uploads everything to `{TARGET_REPO}`.\n\n"
+        "**Note:** Because we are using a free Space (2 vCPUs, 16GB RAM), this script is designed to process one file at a time "
+        "and aggressively clear memory/disk caches. It will take some time, but it won't crash."
     )
-    
-    with gr.Row():
-        hf_token_input = gr.Textbox(
-            label="Hugging Face Token (Write Access)", 
-            type="password",
-            placeholder="hf_... (Leave blank if using Space Secrets)"
-        )
-        
-    with gr.Row():
-        source_input = gr.Textbox(
-            label="Source (Repo ID, Bucket ID, or Local Path)",
-            placeholder="e.g. user/model, username/my-bucket, or /mnt/bucket"
-        )
-        source_type_input = gr.Radio(
-            choices=["repo", "bucket", "local"],
-            value="repo",
-            label="Source Type"
-        )
-        
-    with gr.Row():
-        target_input = gr.Textbox(
-            label="Target Repo ID",
-            placeholder="e.g. your-username/new-model"
-        )
-        repo_type_input = gr.Radio(
-            choices=["model", "dataset", "space"], 
-            value="model", 
-            label="Target Repository Type"
-        )
 
     with gr.Row():
-        precision_input = gr.Dropdown(
-            choices=["None", "FP8", "FP16", "BF16"], 
-            value="None", 
-            label="On-the-Fly Quantization"
-        )
-        target_folders_input = gr.Textbox(
-            label="Target Folders for Quantization (Comma separated)",
-            placeholder="e.g. text_encoder, transformer",
-            value="text_encoder, transformer"
-        )
+        with gr.Column(scale=2):
+            hf_token = gr.Textbox(
+                label="Hugging Face Token (Needs Write Access)", 
+                type="password",
+                placeholder="hf_..."
+            )
+            start_btn = gr.Button("Start Quantization & Upload", variant="primary")
         
-    clone_btn = gr.Button("Stealth Clone & Process", variant="primary")
-    output = gr.Textbox(label="Status", lines=3)
-    
-    clone_btn.click(
-        fn=stealth_clone_hf_repo, 
-        inputs=[
-            hf_token_input, source_input, source_type_input, target_input, 
-            repo_type_input, precision_input, target_folders_input
-        ], 
-        outputs=output
+        with gr.Column(scale=3):
+            output_log = gr.Textbox(
+                label="Operation Logs", 
+                lines=15, 
+                interactive=False,
+                max_lines=20
+            )
+
+    start_btn.click(
+        fn=convert_and_upload, 
+        inputs=[hf_token], 
+        outputs=[output_log]
     )
 
 if __name__ == "__main__":