fix: can't change or delete preview (#100)

py/download.py

@@ -180,8 +180,8 @@ async def create_model_download_task(task_data: dict, request):
         raise RuntimeError(f"Task {task_id} already exists")
 
     try:
-        preview_url = task_data.pop("preview", None)
-        utils.save_model_preview_image(task_path, preview_url)
+        previewFile = task_data.pop("previewFile", None)
+        utils.save_model_preview_image(task_path, previewFile)
         set_task_content(task_id, task_data)
         task_status = TaskStatus(
             taskId=task_id,
@@ -361,9 +361,7 @@ async def download_model_file(
     )
 
     if response.status_code not in (200, 206):
-        raise RuntimeError(
-            f"Failed to download {task_content.fullname}, status code: {response.status_code}"
-        )
+        raise RuntimeError(f"Failed to download {task_content.fullname}, status code: {response.status_code}")
 
     # Some models require logging in before they can be downloaded.
     # If no token is carried, it will be redirected to the login page.
@@ -376,9 +374,7 @@ async def download_model_file(
         # If it cannot be downloaded, a redirect will definitely occur.
         # Maybe consider getting the redirect url from response.history to make a judgment.
         # Here we also need to consider how different websites are processed.
-        raise RuntimeError(
-            f"{task_content.fullname} needs to be logged in to download. Please set the API-Key first."
-        )
+        raise RuntimeError(f"{task_content.fullname} needs to be logged in to download. Please set the API-Key first.")
 
     # When parsing model information from HuggingFace API,
     # the file size was not found and needs to be obtained from the response header.

py/services.py

@@ -73,7 +73,10 @@ def update_model(model_path: str, model_data: dict):
 
     if "previewFile" in model_data:
         previewFile = model_data["previewFile"]
-        utils.save_model_preview_image(model_path, previewFile)
+        if type(previewFile) is str and previewFile == "undefined":
+            utils.remove_model_preview_image(model_path)
+        else:
+            utils.save_model_preview_image(model_path, previewFile)
 
     if "description" in model_data:
         description = model_data["description"]

py/utils.py

@@ -249,19 +249,45 @@ from PIL import Image
 from io import BytesIO
 
 
-def save_model_preview_image(model_path: str, image_url: str):
-    try:
-        image_response = requests.get(image_url)
-        image_response.raise_for_status()
+def remove_model_preview_image(model_path: str):
+    basename = os.path.splitext(model_path)[0]
+    preview_path = f"{basename}.webp"
+    if os.path.exists(preview_path):
+        os.remove(preview_path)
 
-        basename = os.path.splitext(model_path)[0]
-        preview_path = f"{basename}.webp"
-        image = Image.open(BytesIO(image_response.content))
+
+def save_model_preview_image(model_path: str, image_file_or_url: Any):
+    basename = os.path.splitext(model_path)[0]
+    preview_path = f"{basename}.webp"
+
+    # Download image file if it is url
+    if type(image_file_or_url) is str:
+        image_url = image_file_or_url
+
+        try:
+            image_response = requests.get(image_url)
+            image_response.raise_for_status()
+
+            image = Image.open(BytesIO(image_response.content))
+            image.save(preview_path, "WEBP")
+
+        except Exception as e:
+            print_error(f"Failed to download image: {e}")
+
+    else:
+        # Assert image as file
+        image_file = image_file_or_url
+
+        if not isinstance(image_file, web.FileField):
+            raise RuntimeError("Invalid image file")
+
+        content_type: str = image_file.content_type
+        if not content_type.startswith("image/"):
+            raise RuntimeError(f"FileTypeError: expected image, got {content_type}")
+
+        image = Image.open(image_file.file)
         image.save(preview_path, "WEBP")
 
-    except Exception as e:
-        print_error(f"Failed to download image: {e}")
-
 
 def get_model_all_descriptions(model_path: str):
     base_dirname = os.path.dirname(model_path)