Add flags.

app.py

@@ -16,6 +16,16 @@ api = HfApi()
 def restart_space():
     api.restart_space(repo_id="civitaiarchive/civitai-to-hf-uploader", token=os.environ["HF_TOKEN"])
 
+BOOLPARM = {
+    # Raw label conversion.
+    "Default": None,
+    "Include": True,
+    "Exclude": False,
+    # Url parm conversion.
+    True: "true",
+    False: "false",
+    None: "", # Should not pass parm in this case.
+}
 
 def download_file(url, file_path, folder, api_key=None):
     headers = {}
@@ -43,10 +53,15 @@ def download_file(url, file_path, folder, api_key=None):
         raise gr.Error(f"Error downloading file: {str(e)}")
 
         
-def get_files_by_username(username, api_key=None):
-    url = f"https://civitai.com/api/v1/models?username={username}&limit=100&nsfw=true"
+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 = {}
     headers = {}
+    if nsfw is not None:
+        url = url + f"&nsfw={BOOLPARM[nsfw]}"
+    if hidden is not None:
+        url = url + f"&hidden={BOOLPARM[hidden]}"
+    gr.Info(f"SBM url: {url}")
     if api_key:
         headers['Authorization'] = f'Bearer {api_key}'
 
@@ -97,13 +112,13 @@ def get_files_by_model_id(model_id, api_key=None):
     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):
+def process_url(url, profile, user_repo_id, oauth_token, folder, api_key=None, nsfw=None, hidden=None):
     if url.startswith("https://civitai.com/models/"):
         model_id = url.split('/')[4]
         files = 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)
+        files = 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")
     
@@ -175,7 +190,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):
+def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuthToken, url, destination_repo, civitai_api_key=None, nsfw=None, hidden=None):
     if not profile.name:
         raise gr.Error("Are you sure you are logged in?")
     
@@ -186,6 +201,10 @@ def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], oauth_token: gr.OAuth
     if not re.match(r'^[a-zA-Z0-9_-]+$', destination_repo):
         raise gr.Error("Destination repository name must contain only alphanumeric characters, underscores, and hyphens")
 
+    # Convert bool flags.
+    nsfw = BOOLPARM[nsfw]
+    hidden = BOOLPARM[hidden]
+
     folder = str(uuid.uuid4())
     os.makedirs(folder, exist_ok=False)
     gr.Info(f"Starting download from {url}")
@@ -201,7 +220,7 @@ 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)
+        files = process_url(url, profile, user_repo_id, oauth_token.token, folder, civitai_api_key, nsfw, hidden)
         if not files or len(files.keys()) == 0:
             raise gr.Error("No files were copied. Something went wrong.")
         
@@ -268,6 +287,16 @@ Once uploaded, it will add this repository to CivitaiArchive.com as a mirror.
             label="CivitAI API Key",
             info="Optional: Provide your own CivitAI API key to avoid rate limits. If not provided, a default key will be used.",
         )
+        include_nsfw = gr.Dropdown(
+            choices=["Default", "Include", "Exclude"],
+            value="Include",
+            label="Optional: Include, exclude or do not specify inclusion of nsfw models."
+        )
+        include_hidden = gr.Dropdown(
+            choices=["Default", "Include", "Exclude"],
+            value="Default",
+            label="Optional: Include, exclude or do not specify inclusion of hidden models."
+        )
         
     instructions = gr.HTML("")
     submit_button_civit = gr.Button("Upload to Hugging Face", interactive=True)
@@ -277,7 +306,7 @@ Once uploaded, it will add this repository to CivitaiArchive.com as a mirror.
     
     submit_button_civit.click(
         fn=upload_civit_to_hf,
-        inputs=[submit_source_civit, destination_repo, civitai_api_key],
+        inputs=[submit_source_civit, destination_repo, civitai_api_key, include_nsfw, include_hidden],
         outputs=[output]
     )