ford442 commited on
Commit
700ab04
·
verified ·
1 Parent(s): 7681839

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -0
app.py CHANGED
@@ -111,6 +111,11 @@ def load_and_prepare_model():
111
  #pipe.enable_vae_tiling()
112
  pipe.to(device=device, dtype=torch.bfloat16)
113
  pipe.unet.set_attn_processor(AttnProcessor2_0())
 
 
 
 
 
114
  return pipe
115
 
116
  pipe = load_and_prepare_model()
 
111
  #pipe.enable_vae_tiling()
112
  pipe.to(device=device, dtype=torch.bfloat16)
113
  pipe.unet.set_attn_processor(AttnProcessor2_0())
114
+ for name, param in pipe.unet.named_parameters():
115
+ if 'qkv' in name or 'norm' in name or 'proj' in name:
116
+ param.requires_grad_()
117
+ else:
118
+ param.requires_grad_(False)
119
  return pipe
120
 
121
  pipe = load_and_prepare_model()