Update app.py

app.py

@@ -59,15 +59,19 @@ class ZeroGPUAuraSR(AuraSR):
 
 
 aura_sr = ZeroGPUAuraSR.from_pretrained("fal/AuraSR-v2")
+aura_sr_v1 = ZeroGPUAuraSR.from_pretrained("fal-ai/AuraSR")
 
 
 @spaces.GPU()
-def predict(img):
-    return aura_sr.upscale_4x(img)
+def predict(img, model_selection):
+    return {'v1': aura_sr_v1, 'v2': aura_sr}.get().upscale_4x(img)
 
 
-demo = gr.Interface(predict, inputs=gr.Image(), 
-    outputs=gr.Image())
+demo = gr.Interface(
+    predict,
+    inputs=[gr.Image(), gr.Dropdown(value='v2', choices=['v1', 'v2'])],
+    outputs=gr.Image()
+)
 
 
 demo.launch()