Spaces:
Running
on
Zero
Running
on
Zero
add some logs
Browse files
Rodin.py
CHANGED
@@ -20,12 +20,12 @@ def login(email, password):
|
|
20 |
try:
|
21 |
response_data = response.json()
|
22 |
except json.JSONDecodeError as e:
|
23 |
-
|
24 |
raise e
|
25 |
|
26 |
if 'error' in response_data and response_data['error']:
|
27 |
raise Exception(response_data['error'])
|
28 |
-
|
29 |
user_uuid = response_data['user_uuid']
|
30 |
token = response_data['token']
|
31 |
|
@@ -122,6 +122,9 @@ def load_image(img_path):
|
|
122 |
image_bytes = byte_io.getvalue()
|
123 |
return image_bytes
|
124 |
|
|
|
|
|
|
|
125 |
class Generator:
|
126 |
def __init__(self, user_id, password) -> None:
|
127 |
_, self.token = login(user_id, password)
|
@@ -131,23 +134,23 @@ class Generator:
|
|
131 |
|
132 |
def preprocess(self, prompt, image_path, cache_image_base64, task_uuid=""):
|
133 |
if cache_image_base64 and prompt and task_uuid != "":
|
134 |
-
|
135 |
return prompt, cache_image_base64
|
136 |
-
|
137 |
success = False
|
138 |
while not success:
|
139 |
-
|
140 |
image_file = load_image(image_path)
|
141 |
if prompt and task_uuid:
|
142 |
preprocess_response = rodin_preprocess_image(generate_prompt=False, image=image_file, name="images.png", token=self.token)
|
143 |
else:
|
144 |
preprocess_response = rodin_preprocess_image(generate_prompt=True, image=image_file, name="images.png", token=self.token)
|
145 |
-
|
146 |
if 'error' in preprocess_response:
|
147 |
-
|
148 |
raise RuntimeError
|
149 |
elif preprocess_response.get("statusCode") == 401:
|
150 |
-
|
151 |
_, self.token = login(self.user_id, self.password)
|
152 |
continue
|
153 |
else:
|
@@ -157,13 +160,13 @@ class Generator:
|
|
157 |
processed_image = "data:image/png;base64," + preprocess_response.get('processed_image', None)
|
158 |
success = True
|
159 |
except Exception as e:
|
160 |
-
|
161 |
raise e
|
162 |
|
163 |
return prompt, processed_image
|
164 |
|
165 |
def generate_mesh(self, prompt, processed_image, task_uuid=""):
|
166 |
-
|
167 |
if task_uuid == "":
|
168 |
settings = {'view_weights': [1]} # Define weights as per your requirements, for multiple images, use multiple values, e,g [0.5, 0.5]
|
169 |
images = [processed_image] # List of images, all the images should be processed first
|
@@ -173,7 +176,7 @@ class Generator:
|
|
173 |
try:
|
174 |
progress_checker.start()
|
175 |
except Exception as e:
|
176 |
-
|
177 |
time.sleep(5)
|
178 |
|
179 |
task_uuid = mesh_response['uuid'] # The task_uuid should be same during whole generation process
|
@@ -193,18 +196,17 @@ class Generator:
|
|
193 |
try:
|
194 |
checker.start()
|
195 |
except Exception as e:
|
196 |
-
|
197 |
time.sleep(5)
|
198 |
|
199 |
preview_image = next(reversed(rodin_history(task_uuid, self.token).items()))[1]["preview_image"]
|
200 |
-
# print(f"Preview image URL: {rodin_history(task_uuid, self.token)}")
|
201 |
response = requests.get(preview_image, stream=True)
|
202 |
if response.status_code == 200:
|
203 |
# 创建一个PIL Image对象
|
204 |
image = Image.open(response.raw)
|
205 |
# 在这里对image对象进行处理,如显示、保存等
|
206 |
else:
|
207 |
-
|
208 |
raise RuntimeError
|
209 |
response.close()
|
210 |
return image, task_uuid, crop_image(image, DEFAULT)
|
|
|
20 |
try:
|
21 |
response_data = response.json()
|
22 |
except json.JSONDecodeError as e:
|
23 |
+
log("ERROR", f"Error in login: {response}")
|
24 |
raise e
|
25 |
|
26 |
if 'error' in response_data and response_data['error']:
|
27 |
raise Exception(response_data['error'])
|
28 |
+
log("INFO", f"Logged successfully")
|
29 |
user_uuid = response_data['user_uuid']
|
30 |
token = response_data['token']
|
31 |
|
|
|
122 |
image_bytes = byte_io.getvalue()
|
123 |
return image_bytes
|
124 |
|
125 |
+
def log(level, info_text):
|
126 |
+
print(f"[ {level} ] - {time.strftime("%Y%m%d_%H%M%S", time.localtime())} - {info_text}")
|
127 |
+
|
128 |
class Generator:
|
129 |
def __init__(self, user_id, password) -> None:
|
130 |
_, self.token = login(user_id, password)
|
|
|
134 |
|
135 |
def preprocess(self, prompt, image_path, cache_image_base64, task_uuid=""):
|
136 |
if cache_image_base64 and prompt and task_uuid != "":
|
137 |
+
log("INFO", "Using cached image and prompt...")
|
138 |
return prompt, cache_image_base64
|
139 |
+
log("INFO", "Preprocessing image...")
|
140 |
success = False
|
141 |
while not success:
|
142 |
+
log("INFO", "Loading image...")
|
143 |
image_file = load_image(image_path)
|
144 |
if prompt and task_uuid:
|
145 |
preprocess_response = rodin_preprocess_image(generate_prompt=False, image=image_file, name="images.png", token=self.token)
|
146 |
else:
|
147 |
preprocess_response = rodin_preprocess_image(generate_prompt=True, image=image_file, name="images.png", token=self.token)
|
148 |
+
log("INFO", f"Image preprocessed: {preprocess_response.get('statusCode')}")
|
149 |
if 'error' in preprocess_response:
|
150 |
+
log("ERROR", f"Error in image preprocessing: {preprocess_response['error']}")
|
151 |
raise RuntimeError
|
152 |
elif preprocess_response.get("statusCode") == 401:
|
153 |
+
log("WARNING", "Token expired. Logging in again...")
|
154 |
_, self.token = login(self.user_id, self.password)
|
155 |
continue
|
156 |
else:
|
|
|
160 |
processed_image = "data:image/png;base64," + preprocess_response.get('processed_image', None)
|
161 |
success = True
|
162 |
except Exception as e:
|
163 |
+
log("ERROR", f"Error in image preprocessing: {preprocess_response}")
|
164 |
raise e
|
165 |
|
166 |
return prompt, processed_image
|
167 |
|
168 |
def generate_mesh(self, prompt, processed_image, task_uuid=""):
|
169 |
+
log("INFO", "Generating mesh...")
|
170 |
if task_uuid == "":
|
171 |
settings = {'view_weights': [1]} # Define weights as per your requirements, for multiple images, use multiple values, e,g [0.5, 0.5]
|
172 |
images = [processed_image] # List of images, all the images should be processed first
|
|
|
176 |
try:
|
177 |
progress_checker.start()
|
178 |
except Exception as e:
|
179 |
+
log("ERROR", f"Error in generating mesh: {e}")
|
180 |
time.sleep(5)
|
181 |
|
182 |
task_uuid = mesh_response['uuid'] # The task_uuid should be same during whole generation process
|
|
|
196 |
try:
|
197 |
checker.start()
|
198 |
except Exception as e:
|
199 |
+
log("ERROR", f"Error in updating mesh: {e}")
|
200 |
time.sleep(5)
|
201 |
|
202 |
preview_image = next(reversed(rodin_history(task_uuid, self.token).items()))[1]["preview_image"]
|
|
|
203 |
response = requests.get(preview_image, stream=True)
|
204 |
if response.status_code == 200:
|
205 |
# 创建一个PIL Image对象
|
206 |
image = Image.open(response.raw)
|
207 |
# 在这里对image对象进行处理,如显示、保存等
|
208 |
else:
|
209 |
+
log("ERROR", f"Error in generating mesh: {response}")
|
210 |
raise RuntimeError
|
211 |
response.close()
|
212 |
return image, task_uuid, crop_image(image, DEFAULT)
|