Update app.py

app.py

@@ -158,6 +158,23 @@ def stage1_process(
     print('<<== stage1_process')
     return LQ, gr.update(visible = True)
 
+from torch.cuda.amp import autocast  # autocastをインポート
+
+@spaces.GPU(duration=20)  # GPUを利用する関数にデコレーターを追加
+@torch.no_grad()
+def llave_process(input_image, temperature, top_p, qs=None):
+    torch.cuda.set_device(LLaVA_device)
+    with autocast():  # AMPを使用
+        if use_llava and llava_agent is not None:
+            LQ = HWC3(input_image)
+            LQ = Image.fromarray(LQ.astype('uint8'))
+            captions = llava_agent.gen_image_caption([LQ], temperature=temperature, top_p=top_p, qs=qs)
+        else:
+            captions = ['LLaVA is not available. Please add text manually.']
+    torch.cuda.empty_cache()  # メモリを解放
+    return captions[0]
+
+
 def stage2_process(*args, **kwargs):
     try:
         return restore_in_Xmin(*args, **kwargs)