mskrt commited on
Commit
6a798c0
verified
1 Parent(s): cf1bda4

Upload 13 files

Browse files
Files changed (1) hide show
  1. pipeline.py +3 -2
pipeline.py CHANGED
@@ -43,6 +43,7 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
43
  #self.scheduler = scheduler
44
 
45
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
46
  vae.to(device)
47
  unet.to(device)
48
  text_encoder.to(device)
@@ -51,8 +52,8 @@ class SuperDiffPipeline(DiffusionPipeline, ConfigMixin):
51
  vae=vae,
52
  text_encoder=text_encoder,
53
  tokenizer=tokenizer,
54
- **kwargs
55
- )
56
 
57
  #self.register_to_config(
58
  # vae=vae.__class__.__name__,
 
43
  #self.scheduler = scheduler
44
 
45
  device = "cuda" if torch.cuda.is_available() else "cpu"
46
+
47
  vae.to(device)
48
  unet.to(device)
49
  text_encoder.to(device)
 
52
  vae=vae,
53
  text_encoder=text_encoder,
54
  tokenizer=tokenizer,
55
+ **kwargs)
56
+
57
 
58
  #self.register_to_config(
59
  # vae=vae.__class__.__name__,