amosyou commited on
Commit
2fa3fb0
β€’
1 Parent(s): a4e01fa

fix: move to cuda before generate

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -65,6 +65,7 @@ model.load_state_dict(torch.load(
65
  model.eval()
66
  model.ingrs_only = False
67
  model.recipe_only = False
 
68
 
69
  transform_list = []
70
  transform_list.append(transforms.ToTensor())
@@ -80,13 +81,12 @@ temperature = 1.0
80
  numgens = 1
81
 
82
  # StableDiffusion
83
- pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4')
84
 
85
  @spaces.GPU
86
  def generate_image(input_img):
87
 
88
  # Inverse Cooking
89
- model = model.to(device)
90
  image_tensor = transform(input_img).unsqueeze(0).to(device)
91
 
92
  for i in range(numgens):
@@ -110,7 +110,6 @@ def generate_image(input_img):
110
  # {"prompt": prompt, "ingredients": ingredients, "ingr_ids": ingr_ids}
111
 
112
  # StableDiffusion
113
- pipe = pipe.to(device)
114
  new_image = pipe(prompt).images[0]
115
  return new_image
116
 
 
65
  model.eval()
66
  model.ingrs_only = False
67
  model.recipe_only = False
68
+ model = model.to(device)
69
 
70
  transform_list = []
71
  transform_list.append(transforms.ToTensor())
 
81
  numgens = 1
82
 
83
  # StableDiffusion
84
+ pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4').to(device)
85
 
86
  @spaces.GPU
87
  def generate_image(input_img):
88
 
89
  # Inverse Cooking
 
90
  image_tensor = transform(input_img).unsqueeze(0).to(device)
91
 
92
  for i in range(numgens):
 
110
  # {"prompt": prompt, "ingredients": ingredients, "ingr_ids": ingr_ids}
111
 
112
  # StableDiffusion
 
113
  new_image = pipe(prompt).images[0]
114
  return new_image
115