fix: move to cuda in generate function

app.py

@@ -59,7 +59,6 @@ model = get_model(args, ingr_vocab_size, instrs_vocab_size)
 model.load_state_dict(torch.load(
     hf_hub_download(REPO_ID, 'data/modelbest.ckpt', token=HF_TOKEN), map_location=map_loc)
 )
-model = model.to(device)
 model.eval()
 model.ingrs_only = False
 model.recipe_only = False
@@ -78,12 +77,13 @@ temperature = 1.0
 numgens = 1
 
 # StableDiffusion
-pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4').to('cuda')
+pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')
 
 @spaces.GPU
 def generate_image(input_img):
 
     # Inverse Cooking
+    model = model.to(device)
     image_tensor = transform(input_img).unsqueeze(0).to(device)
 
     for i in range(numgens):
@@ -107,6 +107,7 @@ def generate_image(input_img):
     # {"prompt": prompt, "ingredients": ingredients, "ingr_ids": ingr_ids}
     
     # StableDiffusion
+    pipe = pipe.to(device)
     new_image = pipe(prompt).images[0]
     return new_image