app.py: added controlnet configs.

app.py

@@ -18,7 +18,7 @@ from style_template import styles
 MAX_SEED = np.iinfo(np.int32).max
 STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = 'Watercolor'
-DEFAULT_MODEL_NAME = 'sd_xl_base_1.0'
+DEFAULT_MODEL_NAME = 'sdxlUnstableDiffusers_v8HEAVENSWRATH_133813'
 enable_lcm_arg = False
 
 # Path to InstantID models
@@ -100,6 +100,25 @@ SDXL_MODELS = [
 	"zavychromaxl_v21_129006",
 ]
 
+CONTROLNET_DICT = dict(
+	pose={
+		'model_name': 'controlnet-openpose-sdxl-1.0',
+		'strength': 1,
+		'preprocessor': 'openpose',
+	},
+	depth={
+		'model_name': 'controlnet-depth-sdxl-1.0',
+		'strength': 1,
+		'preprocessor': 'depth',
+	},
+	canny={
+		'model_name': 'controlnet-canny-sdxl-1.0',
+		'strength': 1,
+		'preprocessor': 'canny',
+	},
+)
+
+
 
 def get_novita_client (novita_key):
 	client = NovitaClient(novita_key, os.getenv('NOVITA_API_URI', None))
@@ -263,6 +282,10 @@ def generate_image (
 		ref_image = PIL.Image.open(ref_image_path)
 		width, height = ref_image.size
 
+		CONTROLNET_DICT['pose']['strength'] = pose_strength
+		CONTROLNET_DICT['canny']['strength'] = canny_strength
+		CONTROLNET_DICT['depth']['strength'] = depth_strength
+
 		res = client._post('/v3/async/instant-id', {
 			'extra': {
 				'response_image_type': 'jpeg',
@@ -273,7 +296,7 @@ def generate_image (
 			'prompt': prompt,
 			'negative_prompt': negative_prompt,
 			'controlnet': {
-				'units': [], # TODO
+				'units': [CONTROLNET_DICT[name] for name in controlnet_selection if name in CONTROLNET_DICT],
 			},
 			'image_num': 1,
 			'steps': num_steps,
@@ -290,9 +313,9 @@ def generate_image (
 		def progress (x):
 			print('progress:', x.task.status)
 		final_res = client.wait_for_task_v3(res['task_id'], callback=progress)
+		if final_res is None or final_res.task.status == V3TaskResponseStatus.TASK_STATUS_FAILED:
+			raise RuntimeError(f'Novita task failed: {final_res and final_res.task.status}')
 		print('status:', final_res.task.status)
-		if final_res.task.status == V3TaskResponseStatus.TASK_STATUS_FAILED:
-			raise RuntimeError(f'Novita task failed: {final_res.task.status}')
 
 		final_res.download_images()
 	except Exception as e: