@spaces.GPU(duration=120)

gradio_demo.py

@@ -67,7 +67,7 @@ if torch.cuda.device_count() > 0:
     else:
         llava_agent = None
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def stage1_process(input_image, gamma_correction):
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
@@ -87,7 +87,7 @@ def stage1_process(input_image, gamma_correction):
     LQ = LQ.round().clip(0, 255).astype(np.uint8)
     return LQ
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def llave_process(input_image, temperature, top_p, qs=None):
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')
@@ -101,7 +101,7 @@ def llave_process(input_image, temperature, top_p, qs=None):
         captions = ['LLaVA is not available. Please add text manually.']
     return captions[0]
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale, edm_steps, s_stage1, s_stage2,
                    s_cfg, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction,
                    linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select):
@@ -165,7 +165,7 @@ def stage2_process(input_image, prompt, a_prompt, n_prompt, num_samples, upscale
             Image.fromarray(result).save(f'./history/{event_id[:5]}/{event_id[5:]}/HQ_{i}.png')
     return [input_image] + results, event_id, 3, ''
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def load_and_reset(param_setting):
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')