Spaces:

anforsm
/

movie-diffusion

Paused

Anton Forsman commited on Nov 27, 2023

Commit

60e1bdf

1 Parent(s): 5373453

fixes

Files changed (1) hide show

inference.py CHANGED Viewed

@@ -10,7 +10,6 @@ import io
 from model import Unet, ConditionalUnet, GaussianDiffusion, DiffusionImageAPI
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-device = "cpu"
 def inference1():
   # new image from web page
@@ -21,7 +20,7 @@ def inference():
   model = Unet(
     image_channels=3,
   )
-  model.load_state_dict(torch.load("./model_final.pt"), map_location=torch.device("cpu"))
   diffusion = GaussianDiffusion(
     model=model,

 from model import Unet, ConditionalUnet, GaussianDiffusion, DiffusionImageAPI
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 def inference1():
   # new image from web page
   model = Unet(
     image_channels=3,
   )
+  model.load_state_dict(torch.load("./model_final.pt", map_location=device))
   diffusion = GaussianDiffusion(
     model=model,