Support image upload.

app.py

@@ -52,10 +52,55 @@ def download_file(url, file_path, folder, api_key=None):
     except Exception as e:
         raise gr.Error(f"Error downloading file: {str(e)}")
 
-        
+def get_model_meta(model):
+    """Read model metadata from json (api/models format).
+
+    Includes all versions, files and images therein.
+    """
+    dfiles = dict()
+    dimages = dict()
+    for version in model['modelVersions']:
+        mainfile = ""
+        for file in version['files']:
+            if mainfile == "":
+                mainfile = os.path.splitext(file['name'])[0]
+            dfiles[str(model['id']) + '/' + str(version['id']) + '/' + file['name']] = {
+                'rawModelId': model['id'],
+                'rawVersionId': version['id'],
+                'rawFilename': file['name'],
+                'downloadUrl': file['downloadUrl'],
+                'modelId': model['name'] + ' - ' + version['name'], # Not used?
+                'modelUrl': f"https://civitai.com/models/{model['id']}?modelVersionId={version['id']}",
+                'author': model['creator']['username'],
+                'authorUrl': f"https://civitai.com/user/{model['creator']['username']}",
+                'mirrorUrl': f"https://civitaiarchive.com/models/{model['id']}?modelVersionId={version['id']}",
+            }
+        for image in version['images']:
+            # Image ids can be used for direct access,
+            # but we also need to rename them to the model file to give an indication of attachment.
+            vext = os.path.splitext(image['url'])[-1]
+            imagename = mainfile + '-ID' + str(image['id']) + '.preview' + vext
+            dimages[str(model['id']) + '/' + str(version['id']) + '/' + imagename] = {
+                'rawModelId': model['id'],
+                'rawVersionId': version['id'],
+                'rawFilename': imagename,
+                'downloadUrl': image['url'],
+                'modelId': model['name'] + ' - ' + version['name'], # Not used?
+                'modelUrl': f"https://civitai.com/models/{model['id']}?modelVersionId={version['id']}",
+                'author': model['creator']['username'],
+                'authorUrl': f"https://civitai.com/user/{model['creator']['username']}",
+                'mirrorUrl': f"https://civitaiarchive.com/models/{model['id']}?modelVersionId={version['id']}",
+                # Tags which might be used for filtering in future.
+                'nsfwLevel': image['nsfwLevel'], # 1 is sfw, 4 is suggestive, 8+ is explicit.
+                'poi': image['poi'], # Real person. Content was purged.
+                'type': image['type'], # image or video. We already have the extension.
+            }
+    return (dfiles, dimages)
+
 def get_files_by_username(username, api_key=None, nsfw=None, hidden=None):
     url = f"https://civitai.com/api/v1/models?username={username}&limit=100"
     output = {}
+    images = {}
     headers = {}
     if nsfw is not None:
         url = url + f"&nsfw={BOOLPARM[nsfw]}"
@@ -70,20 +115,13 @@ def get_files_by_username(username, api_key=None, nsfw=None, hidden=None):
         data = response.json()
         # Add current page items to the list
         for model in data['items']:
-            for version in model['modelVersions']:
-                for file in version['files']:
-                    output[str(model['id']) + '/' + str(version['id']) + '/' + file['name']] = {
-                        'downloadUrl': file['downloadUrl'],
-                        'modelId': model['name'] + ' - ' + version['name'],
-                        'modelUrl': f"https://civitai.com/models/{model['id']}?modelVersionId={version['id']}",
-                        'author': model['creator']['username'],
-                        'authorUrl': f"https://civitai.com/user/{model['creator']['username']}",
-                        'mirrorUrl': f"https://civitaiarchive.com/models/{model['id']}?modelVersionId={version['id']}",
-                    }
+            (dfiles, dimages) = get_model_meta(model)
+            output.update(dfiles)
+            images.update(dimages)
                     
         metadata = data.get('metadata', {})
         url = metadata.get('nextPage', None)
-    return output
+    return (output, images)
         
 def get_files_by_model_id(model_id, api_key=None):
     api_url = f"https://civitai.com/api/v1/models/{model_id}"
@@ -97,41 +135,36 @@ def get_files_by_model_id(model_id, api_key=None):
         model = response.json()
 
         output = {}
-        for version in model['modelVersions']:
-            for file in version['files']:
-                output[str(model['id']) + '/' + str(version['id']) + '/' + file['name']] = {
-                        'downloadUrl': file['downloadUrl'],
-                        'modelId': model['name'] + ' - ' + version['name'],
-                        'modelUrl': f"https://civitai.com/models/{model['id']}?modelVersionId={version['id']}",
-                        'author': model['creator']['username'],
-                        'authorUrl': f"https://civitai.com/user/{model['creator']['username']}",
-                        'mirrorUrl': f"https://civitaiarchive.com/models/{model['id']}?modelVersionId={version['id']}",
-                    }
-        return output
+        images = {}
+        (dfiles, dimages) = get_model_meta(model)
+        output.update(dfiles)
+        images.update(dimages)
+        return (output, images)
 
     except requests.exceptions.RequestException as e:
         raise gr.Error("Something went wrong in fetching CivitAI API")
 
-def process_url(url, profile, user_repo_id, oauth_token, folder, api_key=None, nsfw=None, hidden=None):
+def process_url(url, profile, user_repo_id, oauth_token, folder, api_key=None, nsfw=None, hidden=None, indimages = False):
     if url.startswith("https://civitai.com/models/"):
         model_id = url.split('/')[4]
-        files = get_files_by_model_id(model_id, api_key)
+        (files, images) = get_files_by_model_id(model_id, api_key)
     elif url.startswith("https://civitai.com/user/"):
         username = url.split('/')[4]
-        files = get_files_by_username(username, api_key, nsfw, hidden)
+        (files, images) = get_files_by_username(username, api_key, nsfw, hidden)
     else:
         raise gr.Error("Unknown CivitAI URL format, please provide model URL or user profile URL")
+
+    total_files = len(files)
+    total_images = len(images)
     
-    gr.Info(f"Found {len(files)} files to download")
-    
+    gr.Info(f"Found {total_files} files to download, {total_images} images")
 
     downloaded_files = {}
-    total_files = len(files)
     current_file = 1
     for dl_path, data in files.items():
         try:
             download_url = data['downloadUrl']
-            filename = dl_path.split('/')[-1]
+            filename = data["rawFilename"]
             if file_exists(
                 repo_id = user_repo_id,
                 filename = dl_path,
@@ -169,8 +202,50 @@ Mirror: [{data['mirrorUrl']}]({data['mirrorUrl']})
         finally:
             current_file += 1
 
-    return files
+    # Download images if the model has been downloaded.
+    downloaded_images = {}
+    current_image = 1
+    if not indimages:
+        gr.Info(f"Skipping all images.")
+    else:
+        for dl_path, data in images.items():
+            try:
+                download_url = data['downloadUrl']
+                filename = data["rawFilename"]
+                if not file_exists(
+                    repo_id = user_repo_id,
+                    filename = os.path.split(dl_path)[:-1],
+                    token = oauth_token
+                ):
+                    gr.Info(f"Skipping {filename}, model was not created {dl_path}")
+                    continue
+                if file_exists(
+                    repo_id = user_repo_id,
+                    filename = dl_path,
+                    token = oauth_token
+                ):
+                    gr.Info(f"Skipping {filename}, folder exists {dl_path}")
+                    continue
+                # Image download should be no different.
+                gr.Info(f"Downloading {filename} ({current_image}/{total_images})")
+                download_file(download_url, dl_path, folder, api_key)
+                # Upload the image only.
+                gr.Info(f"Uploading {filename} ({current_image}/{total_images})")
+                base_folder = os.path.join(folder, os.path.dirname(dl_path))
+                upload_folder(
+                    folder_path=base_folder,
+                    repo_id=user_repo_id,
+                    repo_type="model",
+                    path_in_repo=os.path.dirname(dl_path),
+                    token=oauth_token
+                )
+                downloaded_images[dl_path] = download_url
+            except Exception as e:
+                gr.Warning(f"Failed to download {dl_path}: {str(e)}")
+            finally:
+                current_image += 1
 
+    return (files, images)
 
 
 def add_mirror(repo_id):
@@ -190,7 +265,7 @@ def add_mirror(repo_id):
 
 
 
-def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuthToken, url, destination_repo, civitai_api_key=None, nsfw=None, hidden=None):
+def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuthToken, url, destination_repo, civitai_api_key=None, nsfw=None, hidden=None, indimages=False):
     if not profile.name:
         raise gr.Error("Are you sure you are logged in?")
     
@@ -220,11 +295,13 @@ def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuth
             gr.Info(f"Repository {user_repo_id} already exists, will update it")
         update_repo_visibility(repo_id=user_repo_id, private=False, token=oauth_token.token)
 
-        files = process_url(url, profile, user_repo_id, oauth_token.token, folder, civitai_api_key, nsfw, hidden)
+        (files, images) = process_url(url, profile, user_repo_id, oauth_token.token, folder, civitai_api_key, nsfw, hidden, indimages)
         if not files or len(files.keys()) == 0:
             raise gr.Error("No files were copied. Something went wrong.")
         
         gr.Info(f"Copied {len(files)} files")
+        if indimages:
+            gr.Info(f"Copied {len(images)} images")
             
         results = []
 
@@ -299,6 +376,11 @@ Once uploaded, it will add this repository to CivitaiArchive.com as a mirror.
             label="Hidden models",
             info="Optional: Include, exclude or do not specify inclusion of hidden models.",
         )
+        upload_images = gr.Checkbox(
+            value=True,
+            label="Upload images",
+            info="Optional: Upload images alongside the models for future access. Beware content restrictions."
+        )
         
     instructions = gr.HTML("")
     submit_button_civit = gr.Button("Upload to Hugging Face", interactive=True)