@spaces.GPU

gradio_demo.py

@@ -13,6 +13,7 @@ from CKPT_PTH import LLAVA_MODEL_PATH
 import einops
 import copy
 import time
+import spaces
 from huggingface_hub import hf_hub_download
 
 hf_hub_download(repo_id="laion/CLIP-ViT-bigG-14-laion2B-39B-b160k", filename="open_clip_pytorch_model.bin", local_dir="laion_CLIP-ViT-bigG-14-laion2B-39B-b160k")
@@ -66,6 +67,7 @@ if torch.cuda.device_count() > 0:
     else:
         llava_agent = None
 
+@spaces.GPU
 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.')
@@ -85,6 +87,7 @@ def stage1_process(input_image, gamma_correction):
     LQ = LQ.round().clip(0, 255).astype(np.uint8)
     return LQ
 
+@spaces.GPU
 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.')
@@ -98,6 +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
 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):
@@ -161,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
 def load_and_reset(param_setting):
     if torch.cuda.device_count() == 0:
         gr.Warning('Set this space to GPU config to make it work.')