model version fix

main.py

@@ -21,6 +21,10 @@ app = FastAPI()
 def generate_image(prompt, model):
     torch.cuda.empty_cache()
 
+    modelArray = model.split(",")
+    modelName = modelArray[0]
+    modelVersion = modelArray[1]
+
     pipeline = StableDiffusionPipeline.from_pretrained(
         str(model), torch_dtype=torch.float16
     )
@@ -36,7 +40,11 @@ def generate_image(prompt, model):
         "assertions": [
             {
                 "label": "com.truepic.custom.ai",
-                "data": {"model_name": model, "model_version": "1.0", "prompt": prompt},
+                "data": {
+                    "model_name": modelName,
+                    "model_version": modelVersion,
+                    "prompt": prompt,
+                },
             }
         ]
     }

static/index.html

@@ -140,13 +140,13 @@
           <div class="custom-select">
             <select id="model">
               <option disabled selected value>Select</option>
-              <option value="runwayml/stable-diffusion-v1-5">
+              <option value="runwayml/stable-diffusion-v1-5,1.5">
                 runwayml/stable-diffusion-v1-5
               </option>
-              <option value="CompVis/stable-diffusion-v1-4">
+              <option value="CompVis/stable-diffusion-v1-4,1.4">
                 CompVis/stable-diffusion-v1-4
               </option>
-              <option value="stabilityai/stable-diffusion-2-1">
+              <option value="stabilityai/stable-diffusion-2-1,2.1">
                 stabilityai/stable-diffusion-2-1
               </option>
             </select>