Update README.md
Browse files
README.md
CHANGED
@@ -51,14 +51,14 @@ gray_inpaintor = AutoModel.from_pretrained(
|
|
51 |
trust_remote_code=True,
|
52 |
)
|
53 |
|
54 |
-
Load the gray2rgb model
|
55 |
gray2rgb = AutoModel.from_pretrained(
|
56 |
'jwengr/stable-diffusion-2-gray-inpaint-to-rgb',
|
57 |
subfolder='gray2rgb',
|
58 |
trust_remote_code=True,
|
59 |
)
|
60 |
|
61 |
-
Move models to GPU
|
62 |
gray_inpaintor.to('cuda')
|
63 |
gray2rgb.to('cuda')
|
64 |
|
|
|
51 |
trust_remote_code=True,
|
52 |
)
|
53 |
|
54 |
+
# Load the gray2rgb model
|
55 |
gray2rgb = AutoModel.from_pretrained(
|
56 |
'jwengr/stable-diffusion-2-gray-inpaint-to-rgb',
|
57 |
subfolder='gray2rgb',
|
58 |
trust_remote_code=True,
|
59 |
)
|
60 |
|
61 |
+
# Move models to GPU
|
62 |
gray_inpaintor.to('cuda')
|
63 |
gray2rgb.to('cuda')
|
64 |
|