fix: move to cuda before generate
Browse files
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 |
|