Skip to content

Commit 76b58df

Browse files
Merge pull request #1 from prajinkhadka/main
Add batch size in the Constructor.
2 parents 7b6edaf + 029c34f commit 76b58df

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

src/dalle2/dalle2.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,9 @@
66
import os
77

88
class Dalle2():
9-
def __init__(self, bearer):
9+
def __init__(self, bearer, batch_size):
1010
self.bearer = bearer
11+
self.batch_size = batch_size
1112

1213
def generate(self, promt):
1314
url = "https://labs.openai.com/api/labs/tasks"
@@ -19,7 +20,7 @@ def generate(self, promt):
1920
"task_type": "text2im",
2021
"prompt": {
2122
"caption": promt,
22-
"batch_size": 6,
23+
"batch_size": self.batch_size,
2324
}
2425
}
2526

0 commit comments

Comments
 (0)