Update app.py

app.py

@@ -27,24 +27,37 @@ def get_lcm(prompt):
     print(result)
     return result[0]
 
-def infer(image_in):
+def get_sdxl_lightning(prompt):
+    client = Client("AP123/SDXL-Lightning")
+    result = client.predict(
+    		prompt,	# str  in 'parameter_1' Textbox component
+    		api_name="/predict"
+    )
+    print(result)
+    return result
+
+def infer(image_in, chosen_method):
     caption = get_caption(image_in)
-    img_var = get_lcm(caption)
+    if chosen_method == "LCM" :
+        img_var = get_lcm(caption)
+    elif chosen_method == "SDXL Lightning" :
+        img_var = get_sdxl_lightning(caption)
     return img_var
 
 gr.Interface(
     title = "Supa Fast Image Variation",
-    description = "Get quick image variation from image input, using <a href='https://huggingface.co/vikhyatk/moondream1' target='_blank'>moondream1</a> for caption, and <a href='https://huggingface.co/spaces/latent-consistency/lcm-lora-for-sdxl' target='_blank'>LCM SDXL</a> for image generation",
+    description = "Get quick image variation from image input, using <a href='https://huggingface.co/vikhyatk/moondream1' target='_blank'>moondream1</a> for caption, and <a href='https://huggingface.co/spaces/latent-consistency/lcm-lora-for-sdxl' target='_blank'>LCM SDXL</a> or <a href='https://huggingface.co/spaces/AP123/SDXL-Lightning' target='_blank'>SDXL Lightning</a> for image generation",
     fn = infer,
     inputs = [
-        gr.Image(type="filepath", label="Image input")
+        gr.Image(type="filepath", label="Image input"),
+        gr.Dropdown(label="Choose a model", choices=["LCM", "SDXL Lightning"], value="SDXL Lightning")
     ],
     outputs = [
-        gr.Image(label="LCM Image variation")
+        gr.Image(label="Image variation")
     ],
     examples = [
-        "examples/frog_clean.jpg",
-        "examples/martin_pecheur.jpeg",
-        "examples/forest_deer.png"
+        ["examples/frog_clean.jpg", "SDXL Lightning"],
+        ["examples/martin_pecheur.jpeg", "LCM"],
+        ["examples/forest_deer.png", "SDXL Lightning"]
     ]
 ).queue(max_size=25).launch(show_api=False, show_error=True)