Update app.py

app.py

@@ -34,11 +34,19 @@ def extract_info(json_data):
             if model_version["baseModel"] in ["SDXL 1.0", "SDXL 0.9"]:
                 for file in model_version["files"]:
                     if file["primary"]:
+                        # Start by adding the primary file to the list
+                        urls_to_download = [{"url": file["downloadUrl"], "filename": file["name"], "type": "weightName"}]
+                        
+                        # Then append all image URLs to the list
+                        for image in model_version["images"]:
+                            urls_to_download.append({
+                                "url": image["url"],
+                                "filename": os.path.basename(image["url"]),
+                                "type": "imageName"
+                            })
+                        
                         info = {
-                            "urls_to_download": [
-                                {"url": file["downloadUrl"], "filename": file["name"], "type": "weightName"},
-                                {"url": model_version["images"][0]["url"], "filename": os.path.basename(model_version["images"][0]["url"]), "type": "imageName"}
-                            ],
+                            "urls_to_download": urls_to_download,
                             "id": model_version["id"],
                             "modelId": model_version["modelId"],
                             "name": json_data["name"],
@@ -98,9 +106,9 @@ tags:
   - lora
   - diffusers
 base_model: stabilityai/stable-diffusion-xl-base-1.0
-instance_prompt: {info["trainedWords"][0]}
+instance_prompt: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
 widget:
-  - text: {info["trainedWords"][0]}
+  - text: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
 ---
 
 # {info["name"]}
@@ -110,7 +118,13 @@ widget:
 ![Image]({downloaded_files["imageName"][0]})
 
 {info["description"]}
+
 """
+    for index, image in enumerate(downloaded_files["imageName"]):
+        if(index == 1):
+            content += content + f"## Image examples for the model:\n{![Image]({image})}"
+        elif(index > 1):
+            content += content + f"\n{![Image]({image})}"
     readme_content += content + "\n"
 
     with open(f"{folder}/README.md", "w") as file: