Update app.py

app.py

@@ -124,9 +124,8 @@ pipe.to(device=DEVICE)
 #assert len(output.frames[0]) == 16
 #leave_im_emb.detach().to('cpu')
 
-
 @spaces.GPU(duration=10)
-def generate(in_im_embs):
+def generate_gpu(in_im_embs):
     in_im_embs = in_im_embs.to('cuda').unsqueeze(0).unsqueeze(0)
     #im_embs = torch.cat((torch.zeros(1, 1280, device=DEVICE, dtype=dtype), in_im_embs), 0)
 
@@ -136,7 +135,11 @@ def generate(in_im_embs):
                 output.frames[0][len(output.frames[0])//2], DEVICE, 1, output_hidden_state
             )
     im_emb = im_emb.detach().to('cpu').to(torch.float32)
+    return output, im_emb
 
+def generate(in_im_embs):
+    
+    output, im_emb = generate_gpu(in_im_embs)
     nsfw = maybe_nsfw(output.frames[0][len(output.frames[0])//2])
     
     name = str(uuid.uuid4()).replace("-", "")