Update app.py

app.py

@@ -35,7 +35,7 @@ def check_nsfw(json_data, profile):
 def extract_info(json_data):
     if json_data["type"] == "LORA":
         for model_version in json_data["modelVersions"]:
-            if model_version["baseModel"] in ["SDXL 1.0", "SDXL 0.9"]:
+            if model_version["baseModel"] in ["SDXL 1.0", "SDXL 0.9", "SD 1.5", "SD 1.4", "SD 2.1", "SD 2.0", "SD 2.0 768", "SD 2.1 768"]:
                 for file in model_version["files"]:
                     if file["primary"]:
                         # Start by adding the primary file to the list
@@ -52,10 +52,21 @@ def extract_info(json_data):
                                     "type": "imageName",
                                     "prompt": image["meta"]["prompt"] if image["meta"] is not None and "prompt" in image["meta"] else ""
                                 })
-                        
+                        model_mapping = {
+                            "SDXL 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
+                            "SDXL 0.9": "stabilityai/stable-diffusion-xl-base-1.0",
+                            "SD 1.5": "runwayml/stable-diffusion-v1-5",
+                            "SD 1.4": "CompVis/stable-diffusion-v1-4",
+                            "SD 2.1": "stabilityai/stable-diffusion-2-1-base",
+                            "SD 2.0": "stabilityai/stable-diffusion-2-base",
+                            "SD 2.1 768": "stabilityai/stable-diffusion-2-1",
+                            "SD 2.0 768": "stabilityai/stable-diffusion-2"
+                        }
+                        base_model = model_mapping[model_version["baseModel"]]
                         info = {
                             "urls_to_download": urls_to_download,
                             "id": model_version["id"],
+                            "baseModel": base_model,
                             "modelId": model_version["modelId"],
                             "name": json_data["name"],
                             "description": json_data["description"],
@@ -131,7 +142,7 @@ license: other
 tags:
 - {unpacked_tags}
 
-base_model: stabilityai/stable-diffusion-xl-base-1.0
+base_model: {info["baseModel"]}
 instance_prompt: {info['trainedWords'][0] if 'trainedWords' in info and len(info['trainedWords']) > 0 else ''}
 widget:
 {widget_content}