Update app.py

app.py

@@ -122,7 +122,7 @@ def process_url(url, profile, do_download=True, folder="."):
     else:
         raise gr.Error("Something went wrong in fetching CivitAI API")
 
-def create_readme(info, downloaded_files, link_civit=False, is_author=True, folder="."):
+def create_readme(info, downloaded_files, user_repo_id, link_civit=False, is_author=True, folder="."):
     readme_content = ""
     original_url = f"https://civitai.com/models/{info['modelId']}"
     link_civit_disclaimer = f'([CivitAI]({original_url}))'
@@ -131,6 +131,9 @@ def create_readme(info, downloaded_files, link_civit=False, is_author=True, fold
     civit_tags = [t for t in info["tags"] if t not in default_tags]
     tags = default_tags + civit_tags
     unpacked_tags = "\n- ".join(tags)
+    
+    trained_words = info['trainedWords'] if 'trainedWords' in info and info['trainedWords'] else []
+    formatted_words = ', '.join(f'`{word}`' for word in trained_words)
 
     widget_content = ""
     for index, (prompt, image) in enumerate(zip(downloaded_files["imagePrompt"], downloaded_files["imageName"])):
@@ -162,14 +165,39 @@ widget:
 
 {link_civit_disclaimer if link_civit else ''}
 
+## Model description
+
 {info["description"]}
 
+## Trigger words
+
+You should use {formatted_words} to trigger the image generation.
+
+## Download model
+
+Weights for this model are available in Safetensors format.
+
+[Download](/{user_repo_id}/tree/main) them in the Files & versions tab.
+
+## Use it with diffusers
+
+```py
+from diffusers import AutoPipelineForText2Image
+import torch
+
+pipeline = AutoPipelineForText2Image.from_pretrained('{info["baseModel"]}', torch_dtype=torch.float16).to("cuda")
+pipeline.load_lora_weights("{{user_repo_id}, weight_name='{downloaded_files["weightName"]}')
+image = pipeline('{prompt if prompt else (formatted_words if formatted_words else 'Your custom prompt')}').images[0]
+```
+
+For more details, including weighting, merging and fusing LoRAs, check the [documentation on loading LoRAs in diffusers](https://huggingface.co/docs/diffusers/main/en/using-diffusers/loading_adapters)
+
 """
-    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}\n"
-        elif index > 1:
-            content += f"\n![Image {index}]({image})\n> {prompt}\n"
+    #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}\n"
+    #    elif index > 1:
+    #        content += f"\n![Image {index}]({image})\n> {prompt}\n"
     readme_content += content + "\n"
     print(readme_content)
     with open(f"{folder}/README.md", "w") as file:
@@ -262,7 +290,8 @@ def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], url, link_civit=False
     folder = str(uuid.uuid4())
     os.makedirs(folder, exist_ok=False)
     info, downloaded_files = process_url(url, profile, folder=folder)
-    create_readme(info, downloaded_files, link_civit, folder=folder)
+    user_repo_id = f"{profile.preferred_username}/{slug_name}"
+    create_readme(info, downloaded_files, user_repo_id, link_civit, folder=folder)
     try:
         api = HfApi(token=os.environ["HUGGING_FACE_HUB_TOKEN"])
         username = api.whoami()["name"]
@@ -279,7 +308,7 @@ def upload_civit_to_hf(profile: Optional[gr.OAuthProfile], url, link_civit=False
         raise gr.Error("something went wrong")
         
     transfer_repos = gr.load("multimodalart/transfer_repos", hf_token=os.environ["HUGGING_FACE_HUB_TOKEN"], src="spaces")
-    user_repo_id = f"{profile.preferred_username}/{slug_name}"
+    
     response_code = transfer_repos(repo_id, user_repo_id)
     i = 0
     while response_code != "200":