import io from PIL import Image import datetime import asyncio import aiohttp import os from typing import Optional, Tuple, List import logging import folder_paths as comfy_paths ERROR_MEG = { "MISS_API_KEY": "MISS_API_KEY", "MISS_IMAGES_OR_PROMPT": "MISS_IMAGES_OR_PROMPT", "UNKNOWN_ERROR": "UNKNOWN_ERROR", "RODIN_ERROR": "RODIN_ERROR", "NO_JOBS_FOUND": "NO_JOBS_FOUND", "MISS_SUBSCRIPTION_KEY": "MISS_SUBSCRIPTION_KEY", "MISS_JOB_UUID": "MISS_JOB_UUID", "NO_FILES_FOUND": "NO_FILES_FOUND", "RODIN_POLYGEN_ERROR": "RODIN_POLYGEN_ERROR", "MISS_MODEL_ASSET_ID": "MISS_MODEL_ASSET_ID", } SUPPORTED_3D_EXTENSIONS = [ '.obj', '.glb', '.fbx', '.stl', '.usdz', ] QUALITY_MESH_OPTIONS = { "4K-Quad": (4000, "Quad"), "8K-Quad": (8000, "Quad"), "18K-Quad": (18000, "Quad"), "50K-Quad": (50000, "Quad"), "200K-Quad": (200000, "Quad"), "2K-Triangle": (2000, "Raw"), "20K-Triangle": (20000, "Raw"), "150K-Triangle": (150000, "Raw"), "200K-Triangle": (200000, "Raw"), "500K-Triangle": (500000, "Raw"), "1M-Triangle": (1000000, "Raw") } QUALITY_MESH_DEFAULT = { "Gen-2.5-Minimum": (50000, "Raw"), "Gen-2.5-Extreme-Low": (50000, "Raw"), "Gen-2.5-Low": (50000, "Raw"), "Gen-2.5-Medium": (500000, "Raw"), "Gen-2.5-High": (500000, "Raw"), "Gen-2.5-ExtremeHigh": (1000000, "Raw") } RODIN_API_BASE_URL = "https://api.hyper3d.com" MAX_PARALLEL = 3 def tensor_to_filelike(tensor): """ Converts a PyTorch tensor to a file-like object. Args: - tensor (torch.Tensor): A tensor representing an image of shape (H, W, C) where C is the number of channels (3 for RGB), H is height, and W is width. Returns: - io.BytesIO: A file-like object containing the image data. """ array = tensor.cpu().numpy() array = (array * 255).astype('uint8') image = Image.fromarray(array, 'RGB') img_byte_arr = io.BytesIO() image.save(img_byte_arr, format='PNG') # PNG is used for lossless compression img_byte_arr.seek(0) return img_byte_arr async def submit_generate_job( api_key: str, images: Optional[List[str]], prompt: Optional[str], seed: Optional[int], quality: Optional[str], quality_override: Optional[int], mesh_mode: Optional[str] = 'Raw', geometry_file_format: Optional[str] = 'glb', material: Optional[str] = 'PBR', texture_mode: Optional[str] = None, tier: Optional[str] = None, ta_pose: Optional[bool] = False, hd_texture: Optional[bool] = False, model_early_export: Optional[bool] = False, is_micro: Optional[bool] = False, geometry_instruct_mode: Optional[str] = 'faithful', bbox: Optional[str] = None, height_cm: Optional[int] = None, ) -> Tuple[str, Optional[str], Optional[str]]: """Submit a job job""" if not api_key: return ERROR_MEG["MISS_API_KEY"], None, None # Prepare the request url = f"{RODIN_API_BASE_URL}/api/v2/rodin" headers = { "Authorization": f"Bearer {api_key}" } if images is None and prompt is None: return ERROR_MEG["MISS_IMAGES_OR_PROMPT"], None, None # Prepare form_data form_data = aiohttp.FormData() # Handle images if images is not None: if isinstance(images, list): for img in images: if isinstance(img, str): with open(img, "rb") as f: file_content = f.read() form_data.add_field("images", file_content) else: form_data.add_field("images", tensor_to_filelike(img)) else: # Handle single image case if isinstance(images, str): with open(images, "rb") as f: file_content = f.read() form_data.add_field("images", file_content) else: form_data.add_field("images", tensor_to_filelike(images)) # Handle prompt if prompt is not None and prompt != "": form_data.add_field("prompt", prompt) # Handle seed if seed is not None and seed != 0 and seed != "": form_data.add_field("seed", str(seed)) # Handle quality if quality is not None and quality_override is None: form_data.add_field("quality", quality) # Handle quality_override if quality_override is not None: form_data.add_field("quality_override", str(quality_override)) # Handle mesh_mode if mesh_mode is not None: form_data.add_field("mesh_mode", mesh_mode) # Handle geometry_file_format if geometry_file_format is not None: form_data.add_field("geometry_file_format", geometry_file_format) # Handle material if material is not None: form_data.add_field("material", material) # Handle texture_mode if texture_mode is not None and texture_mode != "Default": form_data.add_field("texture_mode", texture_mode) # Handle tier if tier is not None: form_data.add_field("tier", tier) # Handle ta_pose if ta_pose is not None: form_data.add_field("ta_pose", str(ta_pose).lower()) # Handle hd_texture if hd_texture is not None: form_data.add_field("hd_texture", str(hd_texture).lower()) # Handle model_early_export if model_early_export is not None: form_data.add_field("model_early_export", str(model_early_export).lower()) # Handle is_micro if is_micro is not None: form_data.add_field("is_micro", str(is_micro).lower()) # Handle geometry_instruct_mode if geometry_instruct_mode is not None: form_data.add_field("geometry_instruct_mode", geometry_instruct_mode) # Handle bbox if bbox is not None: # Format as [width, height, length] for API form_data.add_field("bbox_condition", str(bbox)) if height_cm is not None: form_data.add_field("height", str(height_cm)) for field in form_data._fields: print(field) # Post requests async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, data=form_data) as resp: response = await resp.json() print(response) if resp.status in [200, 201]: job_uuid = response.get("uuid", None) subscription_key = response.get("jobs", {}).get("subscription_key", None) if job_uuid is None or subscription_key is None: return ERROR_MEG["RODIN_ERROR"], None, None else: return "Success! Job submitted successfully.", job_uuid, subscription_key else: return ERROR_MEG["RODIN_ERROR"], None, None async def poll_job_status( api_key: str, job_uuid: str, subscription_key: str, ) -> Tuple[str, Optional[str], Optional[str]]: """Poll the job status""" if not api_key: return ERROR_MEG["MISS_API_KEY"], None, None if not subscription_key: return ERROR_MEG["MISS_SUBSCRIPTION_KEY"], None, None if not job_uuid: return ERROR_MEG["MISS_JOB_UUID"], None, None # Prepare the request url = f"{RODIN_API_BASE_URL}/api/v2/status" headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } data = {"subscription_key": subscription_key} # Post requests async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=data) as resp: response = await resp.json() if resp.status in [200, 201]: jobs = response.get("jobs", []) if jobs is None or len(jobs) == 0: return ERROR_MEG["NO_JOBS_FOUND"], None, response # Check status of all jobs all_done = True any_failed = False statuses = [] for job in jobs: job_status = job.get("status", "unknown") statuses.append(job_status) if job_status == "Failed": any_failed = True elif job_status != "Done": all_done = False if all_done: print(f"Task {job_uuid} statuses: {statuses}") print(f"Task {job_uuid} Done!") return f"Success! Task {job_uuid} All jobs have completed successfully.", "done", response elif any_failed: return f"Failed! Task {job_uuid} Jobs have failed. Statuses: {statuses}", "failed", response else: print(f"Task {job_uuid} statuses: {statuses}") return f"Running... Task {job_uuid} Some jobs are still running. Statuses: {statuses}", "running", response else: return ERROR_MEG["UNKNOWN_ERROR"], None, response async def download_results( api_key: str, job_uuid: str, geometry_format: str = 'glb', ) -> Tuple[str, Optional[str], Optional[str]]: """Download the job results""" if not api_key: return ERROR_MEG["MISS_API_KEY"], None, None if not job_uuid: return ERROR_MEG["MISS_JOB_UUID"], None, None # Prepare request url = f"{RODIN_API_BASE_URL}/api/v2/download" headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } data = {"task_uuid": job_uuid} # Post requests async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, json=data) as resp: response = await resp.json() if resp.status in [200, 201]: file_list = response.get("list", []) if file_list is None or len(file_list) == 0: return ERROR_MEG["NO_FILES_FOUND"], None, None else: return await _download_all_file(api_key=api_key, file_list=file_list, geometry_format=geometry_format, task_uuid=job_uuid) else: print(f"❌ Error downloading files: {response}") return ERROR_MEG["RODIN_ERROR"], None, None async def _download_all_file( api_key: str, file_list: list[str], geometry_format: str, task_uuid: str, ) -> Tuple[str, Optional[str], Optional[str]]: """Download the file""" if not api_key: return ERROR_MEG["MISS_API_KEY"], None, None if not file_list: return ERROR_MEG["NO_FILES_FOUND"], None, None try: save_path = os.path.join(comfy_paths.get_output_directory(), datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")) os.makedirs(save_path, exist_ok=True) download_count = 0 preview_model_path = None for file_info in file_list: file_url = file_info.get("url", None) file_name = file_info.get("name", f"file_download_{download_count}").split('/')[-1] file_path = os.path.join(save_path, file_name) format_supports_preview = f'.{geometry_format.lower()}' in SUPPORTED_3D_EXTENSIONS if format_supports_preview and preview_model_path is None: preview_model_path = file_path if not file_url: continue print(f"[ download_files ] Downloading file: {file_path}") async with aiohttp.ClientSession() as session: async with session.get(file_url) as resp: with open(file_path, "wb") as f: f.write(await resp.read()) print(f"Downloaded {file_name}") download_count += 1 return f"Success! Downloaded {download_count} files successfully!", preview_model_path, task_uuid except Exception as e: logging.error(f"❌ Error downloading files: {str(e)}") print(e) return ERROR_MEG["UNKNOWN_ERROR"], None, None async def process_full_generation( api_key: str, images: Optional[List[str]], prompt: Optional[str], tier: str, seed: Optional[int], quality: Optional[str], geometry_file_format: Optional[str], material: Optional[str], texture_mode: Optional[str], quality_override: Optional[int], mesh_mode: Optional[str], ta_pose: Optional[bool], hd_texture: Optional[bool], model_early_export: Optional[bool], is_micro: Optional[bool], geometry_instruct_mode: Optional[str], bbox: Optional[str], height_cm: Optional[int] = None, ) -> Tuple[Optional[str], Optional[str]]: """Full pipeline: submit job, poll status, and download results""" # Step 1: Submit job status_msg, job_uuid, subscription_key = await submit_generate_job( api_key=api_key, images=images, prompt=prompt, seed=seed, quality=quality, quality_override=quality_override, mesh_mode=mesh_mode, geometry_file_format=geometry_file_format, material=material, texture_mode=texture_mode, tier=tier, ta_pose=ta_pose, hd_texture=hd_texture, model_early_export=model_early_export, is_micro=is_micro, geometry_instruct_mode=geometry_instruct_mode, bbox=bbox, height_cm=height_cm, ) if job_uuid is None or subscription_key is None: logging.error(f"Failed to submit job: {status_msg}") print(f"Error: {status_msg}") return None, None print(f"Job submitted successfully! Job UUID: {job_uuid}") # Step 2: Poll for completion max_attempts = 360 # 30 minutes with 5-second intervals (1800 seconds) attempt = 0 while attempt < max_attempts: status_msg, ready_flag, full_response = await poll_job_status(api_key=api_key, job_uuid=job_uuid, subscription_key=subscription_key) print(f"Status: {status_msg}") if ready_flag == "done": # Job completed break elif ready_flag == "failed": logging.error(f"Job failed: {status_msg}") print(f"Error: {status_msg}") return None, None await asyncio.sleep(5) attempt += 1 if attempt >= max_attempts: logging.error("Job polling timed out. Please check the job status.") print("Error: Job polling timed out after 30 minutes. Please check the job status on the Rodin dashboard.") return None, None download_status, model_file, task_uuid = await download_results(api_key=api_key, job_uuid=job_uuid, geometry_format=geometry_file_format) print(f"Download status: {download_status}") if not model_file: logging.error("Failed to download model file") print("Error: Failed to download model file. Please check your internet connection and try again.") return None, None print(f"Model downloaded successfully: {model_file}") return model_file, task_uuid async def submit_polygen_task( api_key: str, asset_id: str, model: str, geometry_file_format: Optional[str], mesh_mode: Optional[str], quality: Optional[str], ) -> Tuple[str, Optional[str], Optional[str]]: """Submit a polygen task to the Rodin API""" if not api_key: return ERROR_MEG["MISS_API_KEY"], None, None # Prepare the request url = f"{RODIN_API_BASE_URL}/api/v2/polygen" headers = { "Authorization": f"Bearer {api_key}", } # Prepare data if not model and not asset_id: return ERROR_MEG["MISS_MODEL_ASSET_ID"], None, None if model and asset_id: return ERROR_MEG["PLEASE_ONLY_MODEL_OR_ASSET_ID"], None, None # Prepare form_data form_data = aiohttp.FormData() if model: if isinstance(model, str): file_ext = os.path.splitext(model)[1].lower() content_type_map = { '.obj': 'model/obj', '.glb': 'model/gltf-binary', '.stl': 'model/stl', '.fbx': 'model/fbx', '.usdz': 'model/vnd.usdz+zip', '.usda': 'model/usda', '.usdc': 'model/usdc', } content_type = content_type_map.get(file_ext, 'application/octet-stream') with open(model, "rb") as f: file_content = f.read() form_data.add_field("model", file_content, filename=os.path.basename(model), content_type=content_type) if asset_id: form_data.add_field("asset_id", asset_id) if geometry_file_format: form_data.add_field("geometry_file_format", geometry_file_format) if mesh_mode: form_data.add_field("meshmesh_mode", mesh_mode) if quality: form_data.add_field("quality", quality) # Send the request async with aiohttp.ClientSession() as session: async with session.post(url, headers=headers, data=form_data) as resp: response = await resp.json() print(response) if resp.status in [200, 201]: job_uuid = response.get("uuid", None) subscription_key = response.get("jobs", {}).get("subscription_key", None) if job_uuid is None or subscription_key is None: return ERROR_MEG["RODIN_POLYGEN_ERROR"], None, None else: return "Success! Polygen task submitted successfully.", job_uuid, subscription_key else: return ERROR_MEG["RODIN_POLYGEN_ERROR"], None, None async def full_polygen_pipeline( api_key: str, asset_id: str, model: str, geometry_file_format: Optional[str], mesh_mode: Optional[str], quality: Optional[str], ) -> Tuple[Optional[str], Optional[str]]: """Full pipeline: submit polygen task, poll status, and download results""" # Step 1: Submit polygen task status_msg, job_uuid, subscription_key = await submit_polygen_task( api_key=api_key, asset_id=asset_id, model=model, geometry_file_format=geometry_file_format, mesh_mode=mesh_mode, quality=quality, ) if job_uuid is None or subscription_key is None: logging.error(f"Failed to submit polygen task: {status_msg}") print(f"Error: {status_msg}") return None, None print(f"Polygen task submitted successfully! UUID: {job_uuid}") # Step 2: Poll for completion max_attempts = 360 # 30 minutes with 5-second intervals (1800 seconds) attempt = 0 while attempt < max_attempts: status_msg, ready_flag, full_response = await poll_job_status(api_key=api_key, job_uuid=job_uuid, subscription_key=subscription_key) print(f"Status: {status_msg}") if ready_flag == "done": # Polygen task completed break elif ready_flag == "failed": logging.error(f"Polygen task failed: {status_msg}") print(f"Error: {status_msg}") return None, None await asyncio.sleep(5) attempt += 1 if attempt >= max_attempts: logging.error("Polygen task polling timed out. Please check the task status.") print("Error: Polygen task polling timed out after 30 minutes. Please check the task status on the Rodin dashboard.") return None, None download_status, model_file, _ = await download_results(api_key=api_key, job_uuid=job_uuid, geometry_format=geometry_file_format) print(f"Download status: {download_status}") if not model_file: logging.error("Failed to download model file") print("Error: Failed to download model file. Please check your internet connection and try again.") return None, None print(f"Model downloaded successfully: {model_file}") return model_file, job_uuid