Update app.py

app.py

@@ -29,10 +29,7 @@ notice = "<p><center>Acknowledgment: This demo is powered by GPU resources suppo
 
 # setup
 def setup():
-    device = ("cuda" if torch.cuda.is_available() else "cpu")
     img_pipe = DDPMPipeline.from_pretrained("yuntian-deng/latex2im_ss_finetunegptneo")
-    img_pipe.to(device)
-    
     model_type = "EleutherAI/gpt-neo-125M"
     #encoder = AutoModel.from_pretrained(model_type).to(device)
     encoder = img_pipe.unet.text_encoder
@@ -47,6 +44,8 @@ def setup():
     eos_id = tokenizer.encode(tokenizer.eos_token)[0]
     @spaces.GPU
     def forward_encoder(latex):
+        device = ("cuda" if torch.cuda.is_available() else "cpu")
+        img_pipe.to(device)
         encoded = tokenizer(latex, return_tensors='pt', truncation=True, max_length=1024)
         input_ids = encoded['input_ids']
         input_ids = torch.cat((input_ids, torch.LongTensor([eos_id,]).unsqueeze(0)), dim=-1)