Update app.py

app.py

@@ -1,27 +1,27 @@
 import gradio as gr
 import torch
 import os
+import shutil
 from tempfile import TemporaryDirectory
 from huggingface_hub import hf_hub_download, HfApi
 from safetensors.torch import save_file, load_file
 from collections import defaultdict
 from typing import Dict, List
 
-# --- Логика, скопированная из оригинального скрипта `convert.py` ---
-# Эти внутренние функции нужны для корректной обработки общих (shared) тензоров.
-# Копируем их, чтобы сделать приложение самодостаточным.
-# Источник: https://github.com/huggingface/safetensors/blob/main/safetensors/torch.py
+# --- Logic copied from the original `convert.py` script ---
+# These internal functions are necessary for correctly handling shared tensors.
+# We copy them here to make the application self-contained.
+# Source: https://github.com/huggingface/safetensors/blob/main/safetensors/torch.py
 
 def _is_complete(storage):
+    # The UserWarning from this line can be ignored; it's expected.
     return storage.size() * storage.element_size() == storage.nbytes()
 
 def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[List[str]]:
     tensors = list(state_dict.values())
-    # Can't handle unpickled storages
     storages = {tensor.storage().data_ptr(): [] for tensor in tensors}
     for name, tensor in state_dict.items():
         storages[tensor.storage().data_ptr()].append(name)
-    # Return only tensors that share storage
     return [names for names in storages.values() if len(names) > 1]
 
 def _remove_duplicate_names(
@@ -32,9 +32,6 @@ def _remove_duplicate_names(
     for shared in shareds:
         complete_names = set([name for name in shared if _is_complete(state_dict[name])])
         if not complete_names:
-            # Fallback for very weird cases.
-            # The model is likely to be incorrect after this
-            # but it will be loadable.
             name = list(shared)[0]
             state_dict[name] = state_dict[name].clone()
             complete_names = {name}
@@ -50,15 +47,14 @@ def check_file_size(sf_filename: str, pt_filename: str):
     sf_size = os.stat(sf_filename).st_size
     pt_size = os.stat(pt_filename).st_size
     if (sf_size - pt_size) / pt_size > 0.01:
-        # Не бросаем ошибку, а возвращаем предупреждение
         return (
-            f"ВНИМАНИЕ: Размер сконвертированного файла ({sf_size} байт) "
-            f"более чем на 1% отличается от оригинала ({pt_size} байт)."
+            f"WARNING: The converted file size ({sf_size} bytes) "
+            f"differs from the original ({pt_size} bytes) by more than 1%."
         )
     return None
 
 def convert_file(pt_filename: str, sf_filename: str, device: str):
-    """Основная функция конвертации одного файла."""
+    """Main function to convert a single file."""
     loaded = torch.load(pt_filename, map_location=device, weights_only=True)
     if "state_dict" in loaded:
         loaded = loaded["state_dict"]
@@ -72,110 +68,116 @@ def convert_file(pt_filename: str, sf_filename: str, device: str):
             del loaded[to_remove]
 
     loaded = {k: v.contiguous() for k, v in loaded.items()}
-
     os.makedirs(os.path.dirname(sf_filename), exist_ok=True)
     save_file(loaded, sf_filename, metadata=metadata)
     
     size_warning = check_file_size(sf_filename, pt_filename)
     
-    # Проверка на корректность
     reloaded = load_file(sf_filename)
     for k in loaded:
         pt_tensor = loaded[k].to("cpu")
         sf_tensor = reloaded[k].to("cpu")
         if not torch.equal(pt_tensor, sf_tensor):
-            raise RuntimeError(f"Тензоры не совпадают для ключа {k}!")
+            raise RuntimeError(f"Tensors do not match for key {k}!")
             
     return size_warning
 
 
-# --- Основная логика Gradio-приложения ---
+# --- Main Gradio App Logic ---
 
 def process_model(model_id: str, revision: str, progress=gr.Progress(track_tqdm=True)):
-    """
-    Скачивает, конвертирует и возвращает пути к файлам `.safetensors`.
-    """
     if not model_id:
-        return None, "Ошибка: ID модели не может быть пустым."
+        return None, "Error: Model ID cannot be empty."
 
-    # 1. Определяем устройство (GPU или CPU)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    log_messages = [f"✅ Обнаружено устройство: {device.upper()}"]
+    log_messages = [f"✅ Detected device: {device.upper()}"]
 
     try:
         api = HfApi()
         info = api.model_info(repo_id=model_id, revision=revision)
         filenames = [s.rfilename for s in info.siblings]
     except Exception as e:
-        return None, f"❌ Ошибка: Не ��далось получить информацию о модели `{model_id}`.\n{e}"
+        return None, f"❌ Error: Failed to get model info for `{model_id}`.\n{e}"
 
-    # Ищем файлы для конвертации
     files_to_convert = [f for f in filenames if f.endswith(".bin") or f.endswith(".ckpt")]
     if not files_to_convert:
-        return None, f"ℹ️ В модели `{model_id}` не найдено файлов `.bin` или `.ckpt` для конвертации."
+        return None, f"ℹ️ No .bin or .ckpt files found in model `{model_id}` for conversion."
 
-    log_messages.append(f"🔍 Найдено {len(files_to_convert)} файлов для конвертации: {', '.join(files_to_convert)}")
+    log_messages.append(f"🔍 Found {len(files_to_convert)} file(s) to convert: {', '.join(files_to_convert)}")
 
-    # Используем временную директорию для чистоты
     with TemporaryDirectory() as temp_dir:
-        converted_files = []
-        for filename in progress.tqdm(files_to_convert, desc="Конвертация файлов"):
+        temp_converted_files = []
+        for filename in progress.tqdm(files_to_convert, desc="Converting files"):
             try:
-                # Скачиваем файл
-                log_messages.append(f"\n🚀 Скачивание `{filename}`...")
+                log_messages.append(f"\n🚀 Downloading `{filename}`...")
                 pt_path = hf_hub_download(
-                    repo_id=model_id,
-                    filename=filename,
-                    revision=revision,
+                    repo_id=model_id, filename=filename, revision=revision,
                     cache_dir=os.path.join(temp_dir, "downloads"),
                 )
                 
-                # Конвертируем
-                log_messages.append(f"🛠️ Конвертация `{filename}`...")
-                sf_filename = os.path.splitext(filename)[0] + ".safetensors"
+                log_messages.append(f"🛠️ Converting `{filename}`...")
+                sf_filename = os.path.splitext(os.path.basename(filename))[0] + ".safetensors"
                 sf_path = os.path.join(temp_dir, "converted", sf_filename)
                 
                 size_warning = convert_file(pt_path, sf_path, device)
                 if size_warning:
                     log_messages.append(f"⚠️ {size_warning}")
 
-                converted_files.append(sf_path)
-                log_messages.append(f"✅ Успешно сконвертировано в `{sf_filename}`")
+                temp_converted_files.append(sf_path)
+                log_messages.append(f"✅ Successfully converted to `{sf_filename}`")
             except Exception as e:
-                log_messages.append(f"❌ Ошибка при обработке файла `{filename}`: {e}")
+                log_messages.append(f"❌ Error processing file `{filename}`: {e}")
                 continue
         
-        if not converted_files:
-            return None, "\n".join(log_messages) + "\n\nНе удалось сконвертировать ни один файл."
+        if not temp_converted_files:
+            return None, "\n".join(log_messages) + "\n\nFailed to convert any files."
             
-        final_message = "\n".join(log_messages) + "\n\n" + "🎉 Все файлы успешно обработаны! Готово к скачиванию."
-        return converted_files, final_message
+        # --- KEY CHANGE ---
+        # Copy files from the temporary directory to a persistent (for Gradio) location
+        # before the directory is deleted.
+        persistent_files = []
+        for temp_path in temp_converted_files:
+            # shutil.copy() creates a new file that won't be deleted
+            persistent_path = shutil.copy(temp_path, ".")
+            persistent_files.append(persistent_path)
+        # --------------------
+        
+        final_message = "\n".join(log_messages) + "\n\n" + "🎉 All files processed successfully! Ready for download."
+        # Return the paths to the persistent files
+        return persistent_files, final_message
 
 
-# --- Создание интерфейса Gradio ---
+# --- Create Gradio Interface ---
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # Конвертер моделей в `.safetensors`
-        Эта утилита преобразует веса моделей PyTorch (`.bin`, `.ckpt`) из репозиториев Hugging Face 
-        в безопасный и быстрый формат `.safetensors`.
+        # Model Converter to `.safetensors`
+        This utility converts PyTorch model weights (`.bin`, `.ckpt`) from Hugging Face repositories 
+        to the safe and fast `.safetensors` format.
         
-        **Как использовать:**
-        1. Введите ID модели с Hugging Face (например, `stabilityai/stable-diffusion-2-1-base`).
-        2. Нажмите кноп��у "Конвертировать".
-        3. Дождитесь завершения процесса и скачайте полученные файлы.
+        **How to use:**
+        1. Enter the Model ID from Hugging Face (e.g., `stabilityai/stable-diffusion-2-1-base`).
+        2. Click the "Convert" button.
+        3. Wait for the process to complete and download the resulting files.
         """
     )
     with gr.Row():
-        model_id = gr.Textbox(label="ID модели на Hugging Face", placeholder="например, runwayml/stable-diffusion-v1-5")
-        revision = gr.Textbox(label="Ревизия (ветка)", value="main")
+        model_id = gr.Textbox(label="Hugging Face Model ID", placeholder="e.g., runwayml/stable-diffusion-v1-5")
+        revision = gr.Textbox(label="Revision (branch)", value="main")
     
-    convert_button = gr.Button("Конвертировать", variant="primary")
+    convert_button = gr.Button("Convert", variant="primary")
     
-    gr.Markdown("### Результат")
-    log_output = gr.Markdown(value="Ожидание запуска...")
-    file_output = gr.File(label="Скачать сконвертированные файлы")
+    gr.Markdown("### Result")
+    log_output = gr.Markdown(value="Waiting for input...")
+    file_output = gr.File(label="Download Converted Files")
+
+    gr.Markdown(
+        "<p style='color:grey;font-size:0.8em;'>"
+        "<b>Note:</b> A `UserWarning: TypedStorage is deprecated` message may appear in the logs. "
+        "This is normal and does not affect the result."
+        "</p>"
+    )
 
     convert_button.click(
         fn=process_model,
@@ -184,4 +186,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch()