Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from huggingface_hub import HfApi, update_repo_visibility
 from slugify import slugify
 import gradio as gr
+import re
 import uuid
 from typing import Optional
 import json
@@ -42,7 +43,8 @@ def extract_info(json_data):
                             urls_to_download.append({
                                 "url": image["url"],
                                 "filename": os.path.basename(image["url"]),
-                                "type": "imageName"
+                                "type": "imageName",
+                                "prompt": image["meta"]["prompt"]
                             })
                         
                         info = {
@@ -52,7 +54,8 @@ def extract_info(json_data):
                             "name": json_data["name"],
                             "description": json_data["description"],
                             "trainedWords": model_version["trainedWords"],
-                            "creator": json_data["creator"]["username"]
+                            "creator": json_data["creator"]["username"],
+                            "tags": json_data["tags"]
                         }
                         return info
     return None
@@ -60,11 +63,15 @@ def extract_info(json_data):
 def download_files(info, folder="."):
     downloaded_files = {
         "imageName": [],
+        "imagePrompt": [],
         "weightName": []
     }
     for item in info["urls_to_download"]:
         download_file(item["url"], item["filename"], folder)
         downloaded_files[item["type"]].append(item["filename"])
+        if(item["type"] == "imageName"):
+            prompt_clean = re.sub(r'<.*?>', '', item["prompt"])
+            downloaded_files["imagePrompt"].append(prompt_clean)
     return downloaded_files
 
 def download_file(url, filename, folder="."):
@@ -98,17 +105,17 @@ def create_readme(info, downloaded_files, is_author=True, folder="."):
     readme_content = ""
     original_url = f"https://civitai.com/models/{info['id']}"
     non_author_disclaimer = f'This model was originally uploaded on [CivitAI]({original_url}), by [{info["creator"]}](https://civitai.com/user/{info["creator"]}/models). The information below was provided by the author on CivitAI:'
+    default_tags = ["text-to-image", "stable-diffusion", "lora", "diffusers"]
+    civit_tags = [t for t in info["tags"] if t not in default_tags]
+    widget_prompts = "\n- text:".join(downloaded_files["imagePrompt"])
+    tags = default_tags + civit_tags
+    unpacked_tags = "\n-".join(tags)
     content = f"""---
 license: other
-tags:
-  - text-to-image
-  - stable-diffusion
-  - lora
-  - diffusers
+tags:{unpacked_tags}
 base_model: stabilityai/stable-diffusion-xl-base-1.0
 instance_prompt: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
-widget:
-  - text: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
+widget:{widget_prompts}
 ---
 
 # {info["name"]}
@@ -120,11 +127,11 @@ widget:
 {info["description"]}
 
 """
-    for index, image in enumerate(downloaded_files["imageName"]):
-        if(index == 1):
-            content += f"## Image examples for the model:\n![Image {index}]({image})"
-        elif(index > 1):
-            content += f"\n![Image {index}]({image})"
+    for index, (image, prompt) in enumerate(zip(downloaded_files["imageName"], downloaded_files["imagePrompt"])):
+    if index == 1:
+        content += f"## Image examples for the model:\n![Image {index}]({image})\n> Prompt: {prompt}"
+    elif index > 1:
+        content += f"\n![Image {index}]({image})\n> Prompt: {prompt}"
     readme_content += content + "\n"
 
     with open(f"{folder}/README.md", "w") as file: