Update app.py

app.py

@@ -46,13 +46,13 @@ USE_ZERO_GPU = True
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
     prompt,
+    model,
     seed,
     randomize_seed,
     width,
     height,
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
-    model='2k',
 ):
     print("Using model:", model)
     if model == "2k":
@@ -83,7 +83,7 @@ def infer(
 
 
 if USE_ZERO_GPU:
-    infer = spaces.GPU(infer)
+    infer = spaces.GPU(infer, duration=360)
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -124,14 +124,14 @@ with gr.Blocks(css=css) as demo:
 
             gr.Markdown("### Setting:")
 
-            # model = gr.Radio(
-            #     label="Model",
-            #     choices=[
-            #         ("2K model", "2k"),
-            #         ("4K model (beta)", "4k"),
-            #     ],
-            #     value="2k",
-            # )
+            model = gr.Radio(
+                label="Model",
+                choices=[
+                    ("2K model", "2k"),
+                    ("4K model (beta)", "4k"),
+                ],
+                value="2k",
+            )
 
             with gr.Row():
                 width = gr.Slider(
@@ -176,7 +176,7 @@ with gr.Blocks(css=css) as demo:
         fn=infer,
         inputs=[
             prompt,
-            # model,
+            model,
             seed,
             randomize_seed,
             width,