j-min commited on
Commit
cdecf9e
Β·
1 Parent(s): 95e9355

add cpu offload / xformers

Browse files
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -45,15 +45,19 @@ class Instance:
45
  def load_model(self, model_id='j-min/IterInpaint-CLEVR'):
46
  pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id)
47
 
48
- if torch.cuda.is_available():
49
- pipe = pipe.to('cuda')
50
-
51
  def dummy(images, **kwargs):
52
  return images, False
53
  pipe.safety_checker = dummy
54
  print("Disabled safety checker")
55
 
56
  print("Loaded model")
 
 
 
 
 
 
 
57
  return pipe
58
 
59
  instance = Instance()
 
45
  def load_model(self, model_id='j-min/IterInpaint-CLEVR'):
46
  pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id)
47
 
 
 
 
48
  def dummy(images, **kwargs):
49
  return images, False
50
  pipe.safety_checker = dummy
51
  print("Disabled safety checker")
52
 
53
  print("Loaded model")
54
+
55
+ # This command loads the individual model components on GPU on-demand. So, we don't
56
+ # need to explicitly call pipe.to("cuda").
57
+ pipe.enable_model_cpu_offload()
58
+
59
+ # xformers
60
+ pipe.enable_xformers_memory_efficient_attention()
61
  return pipe
62
 
63
  instance = Instance()