Update app.py

app.py

@@ -760,8 +760,8 @@ def generate_image(
     if seed == 0:
         seed = int(random.random() * 1000000)
         
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch_device = torch.device(device)
+    assert torch.cuda.is_available()
+    torch_device = torch.device("cuda")
     
     global model, model_zero_init
     if not model_zero_init:
@@ -792,8 +792,9 @@ def generate_image(
         timesteps = timesteps[t_idx:]
         x = t * x + (1.0 - t) * init_image.to(x.dtype)
 
-    inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
-    x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
+    with torch_device:
+        inp = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
+        x = denoise(model, **inp, timesteps=timesteps, guidance=guidance)
     
     # with profile(activities=[ProfilerActivity.CPU],record_shapes=True,profile_memory=True) as prof:
     # print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=20))