Update app.py

app.py

@@ -4,7 +4,7 @@ Train a diffusion model on images.
 import gradio as gr
 import argparse
 from einops import rearrange
-from glide_text2im import dist_util, logger
+#from glide_text2im import dist_util, logger
 from torchvision.utils import make_grid
 from glide_text2im.script_util import (
     model_and_diffusion_defaults,
@@ -30,7 +30,7 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
     
     args = parser.parse_args()
     args_up = parser_up.parse_args()
-    dist_util.setup_dist()
+    #dist_util.setup_dist()
 
     if mode == 'sketch':
         args.mode = 'coco-edge'
@@ -59,21 +59,21 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
 
     if  args.model_path:
         print('loading model')
-        model_ckpt = dist_util.load_state_dict(args.model_path, map_location="cpu")
+        model_ckpt = torch.load(args.model_path, map_location="cpu")
 
         model.load_state_dict(
             model_ckpt   , strict=True )
 
     if  args.sr_model_path:
         print('loading sr model')
-        model_ckpt2 = dist_util.load_state_dict(args.sr_model_path, map_location="cpu")
+        model_ckpt2 = torch.load(args.sr_model_path, map_location="cpu")
 
         model_up.load_state_dict(
             model_ckpt2   , strict=True ) 
 
  
-    model.to(dist_util.dev())
-    model_up.to(dist_util.dev())
+    model.cuda()
+    model_up.cuda()
     model.eval()
     model_up.eval()
  
@@ -122,7 +122,7 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
                 prompt=model_kwargs,
                 batch_size= args.num_samples,
                 guidance_scale=args.sample_c,
-                device=dist_util.dev(),
+                device=torch.device('cuda'),
                 prediction_respacing= str(sample_step),
                 upsample_enabled= False,
                 upsample_temp=0.997,
@@ -142,7 +142,7 @@ def run(image, mode, sample_c=1.3,  num_samples=3, sample_step=100):
                 prompt=model_kwargs,
                 batch_size=args.num_samples,
                 guidance_scale=1,
-                device=dist_util.dev(),
+                device=torch.device('cuda'),
                 prediction_respacing= "fast27",
                 upsample_enabled=True,
                 upsample_temp=0.997,