Spaces:
Running
Running
Update app.py
Browse files
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()
|