skkk commited on
Commit
0f9e669
1 Parent(s): 28bc8db
Files changed (1) hide show
  1. Rodin.py +24 -18
Rodin.py CHANGED
@@ -133,24 +133,30 @@ class Generator:
133
  if cache_image_base64 and prompt and task_uuid != "":
134
  return prompt, cache_image_base64
135
  print("Preprocessing image...")
136
- _, self.token = login(self.user_id, self.password)
137
-
138
- image_file = load_image(image_path)
139
- if prompt and task_uuid:
140
- preprocess_response = rodin_preprocess_image(generate_prompt=False, image=image_file, name="images.png", token=self.token)
141
- else:
142
- preprocess_response = rodin_preprocess_image(generate_prompt=True, image=image_file, name="images.png", token=self.token)
143
-
144
- if 'error' in preprocess_response:
145
- print("Error in image preprocessing:", preprocess_response['error'])
146
- else:
147
- if not (prompt and task_uuid):
148
- prompt = preprocess_response.get('prompt', 'Default prompt if none returned')
149
- try:
150
- processed_image = "data:image/png;base64," + preprocess_response.get('processed_image', None)
151
- except Exception as e:
152
- print(preprocess_response)
153
- raise e
 
 
 
 
 
 
154
 
155
  return prompt, processed_image
156
 
 
133
  if cache_image_base64 and prompt and task_uuid != "":
134
  return prompt, cache_image_base64
135
  print("Preprocessing image...")
136
+ success = False
137
+ while not success:
138
+ image_file = load_image(image_path)
139
+ if prompt and task_uuid:
140
+ preprocess_response = rodin_preprocess_image(generate_prompt=False, image=image_file, name="images.png", token=self.token)
141
+ else:
142
+ preprocess_response = rodin_preprocess_image(generate_prompt=True, image=image_file, name="images.png", token=self.token)
143
+
144
+ if 'error' in preprocess_response:
145
+ print("Error in image preprocessing:", preprocess_response['error'])
146
+ raise RuntimeError
147
+ elif preprocess_response.get("statusCode") == "401":
148
+ print("Token expired, refreshing token...")
149
+ _, self.token = login(self.user_id, self.password)
150
+ continue
151
+ else:
152
+ if not (prompt and task_uuid):
153
+ prompt = preprocess_response.get('prompt', 'Default prompt if none returned')
154
+ try:
155
+ processed_image = "data:image/png;base64," + preprocess_response.get('processed_image', None)
156
+ success = True
157
+ except Exception as e:
158
+ print(preprocess_response)
159
+ raise e
160
 
161
  return prompt, processed_image
162