updates to UI

main.py

@@ -18,11 +18,10 @@ print(f"Is CUDA available: {torch.cuda.is_available()}")
 app = FastAPI()
 
 @app.get("/generate")
-def generate_image(prompt):
+def generate_image(prompt, inference_steps, model):
     print(f"Is CUDA available: {torch.cuda.is_available()}")
 
     #model_id = "CompVis/stable-diffusion-v1-4" #stabilityai/stable-diffusion-2-1
-    model_id = "runwayml/stable-diffusion-v1-5"
 
     # Use the DPMSolverMultistepScheduler (DPM-Solver++) scheduler here instead
     #pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
@@ -31,12 +30,12 @@ def generate_image(prompt):
     #pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     #pipe = pipe.to("cuda")
     
-    pipeline = DiffusionPipeline.from_pretrained(model_id)
+    pipeline = DiffusionPipeline.from_pretrained(model)
     #pipeline = pipeline.to("cuda")
     generator = torch.Generator("cpu").manual_seed(0)
     pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
     #image = pipeline(prompt, generator=generator).images[0]
-    image = pipeline(prompt, generator=generator, num_inference_steps=5).images[0]
+    image = pipeline(prompt, generator=generator, num_inference_steps=inference_steps).images[0]
 
 
     #prompt = "a photo of an astronaut riding a horse on mars"
@@ -53,7 +52,7 @@ def generate_image(prompt):
                     {
                     "label": "com.truepic.custom.ai",
                     "data": {
-                        "model_name": model_id,
+                        "model_name": model,
                         "model_version": "1.0",
                         "prompt": prompt
                     }

static/index.html

@@ -14,13 +14,31 @@
         <h1>Text to Image with Signed Attribution</h1>
         <truepic-display><img src="/output.jpg" /></truepic-display>
         <form class="text-gen-form" style="padding:40px 0;">
-          <label for="text-gen-input">Enter number between 0 and 999: </label>
+          <label for="text-gen-input">Enter prompt: </label>
           <input
             id="text-gen-input"
             type="text"
             value=""
             style="margin:40px 0;"
           />
+
+          <label>Inference Steps: </label>
+
+          <select style="width:20%" id="inference_steps">
+            <option value="5">5 (1m 10s)</option>
+            <option value="10">10</option>
+            <option value="20">20</option>
+            <option value="50">50</option>
+          </select>
+
+          <label>Model: </label>
+
+          <select style="width:20%" id="model">
+            <option value="runwayml/stable-diffusion-v1-5">runwayml/stable-diffusion-v1-5</option>
+            <option value="CompVis/stable-diffusion-v1-4">CompVis/stable-diffusion-v1-4</option>
+            <option value="stabilityai/stable-diffusion-2-1">stabilityai/stable-diffusion-2-1</option>
+          </select>
+
           <button id="text-gen-submit">Submit</button>
         </form>
 

static/script.js

@@ -23,8 +23,8 @@ textGenForm.addEventListener('submit', async (event) => {
 
 */ 
 
-const generateImage = async (text) => {
-    const inferResponse = await fetch(`generate?prompt=${text}`);
+const generateImage = async (text, inference_steps, model) => {
+    const inferResponse = await fetch(`generate?prompt=${text}&inference_steps=${inference_steps}&model=${model}`);
     const inferJson = await inferResponse.json();
     
     return inferJson.output;
@@ -35,9 +35,11 @@ textGenForm.addEventListener('submit', async (event) => {
      event.preventDefault();
   
     const textGenInput = document.getElementById('text-gen-input');
+    const inferenceSteps = document.getElementById('inference_steps');
+    const model = document.getElementById('model');
     
     try {
-      const resp = await generateImage(textGenInput.value);
+      const resp = await generateImage(textGenInput.value, inferenceSteps.value, model.value);
       document.getElementById("redirect-form").submit();
     } catch (err) {
       console.error(err);

static/style.css

@@ -11,8 +11,7 @@ body {
   }
   
   main {
-    max-width: 80rem;
-    text-align: center;
+    max-width: 680rem;
   }
   
   section {
@@ -30,12 +29,18 @@ body {
     margin: 0 auto;
   }
   
+  label {
+    display: inline-block;
+    width: 18%;
+  }
+
   input {
-    width: 100%;
+    width: 80%;
   }
   
   button {
     cursor: pointer;
+    display: block;
   }
   
   .text-gen-output {