436 lines
16 KiB
Python
436 lines
16 KiB
Python
import bpy
|
|
import math
|
|
import threading
|
|
import json
|
|
import requests
|
|
import os
|
|
import time
|
|
import queue
|
|
from bpy_extras.io_utils import ImportHelper
|
|
|
|
_status_queue = queue.Queue()
|
|
_model_info_queue = queue.Queue()
|
|
_log_queue = queue.Queue()
|
|
|
|
class ComfyUIPreferences(bpy.types.AddonPreferences):
|
|
bl_idname = __name__
|
|
|
|
comfyui_workflow_path: bpy.props.StringProperty(
|
|
name="Workflow File",
|
|
subtype='FILE_PATH',
|
|
default=""
|
|
)
|
|
|
|
def draw(self, context):
|
|
layout = self.layout
|
|
layout.prop(self, "comfyui_workflow_path")
|
|
|
|
def find_node_by_class(workflow_data, class_type):
|
|
"""Finds the first node ID matching a specific ComfyUI class."""
|
|
for node_id, node in workflow_data.items():
|
|
if node.get("class_type") == class_type:
|
|
return node_id
|
|
return None
|
|
|
|
def find_node_by_title(workflow_data, target_title):
|
|
"""Finds a node ID based on its custom Title."""
|
|
for node_id, node in workflow_data.items():
|
|
meta = node.get("_meta", {})
|
|
if meta.get("title") == target_title:
|
|
return node_id
|
|
return None
|
|
|
|
# ------------------------------
|
|
# Scene property update functions
|
|
# ------------------------------
|
|
def update_input_preview(self, context):
|
|
scene = context.scene
|
|
path = scene.comfyui_input_image
|
|
scene.comfyui_input_preview_path = path if os.path.exists(path) else ""
|
|
for area in context.screen.areas:
|
|
if area.type == 'VIEW_3D':
|
|
area.tag_redraw()
|
|
|
|
def update_comfyui_input_image(self, context):
|
|
path = self.comfyui_input_image
|
|
if path and os.path.exists(path):
|
|
try:
|
|
img = bpy.data.images.load(path, check_existing=True)
|
|
context.scene.comfyui_preview_image = img
|
|
except Exception as e:
|
|
print(f"Failed to load preview: {e}")
|
|
context.scene.comfyui_preview_image = None
|
|
else:
|
|
context.scene.comfyui_preview_image = None
|
|
|
|
# ------------------------------
|
|
# Scene properties registration
|
|
# ------------------------------
|
|
def register_scene_props():
|
|
bpy.types.Scene.comfyui_server_url = bpy.props.StringProperty(default="http://localhost:8188")
|
|
bpy.types.Scene.comfyui_workflow_path = bpy.props.StringProperty(
|
|
default="",
|
|
subtype='FILE_PATH',
|
|
name="Workflow File"
|
|
)
|
|
|
|
bpy.types.Scene.comfyui_input_image = bpy.props.StringProperty(
|
|
name="Input Image",
|
|
subtype='FILE_PATH',
|
|
update=update_comfyui_input_image
|
|
)
|
|
|
|
bpy.types.Scene.comfyui_preview_image = bpy.props.PointerProperty(
|
|
type=bpy.types.Image,
|
|
name="ComfyUI Input Preview"
|
|
)
|
|
bpy.types.Scene.comfyui_exp_warning = bpy.props.BoolProperty(name="Expand Warning", default=True)
|
|
bpy.types.Scene.comfyui_exp_log = bpy.props.BoolProperty(name="Expand Log", default=True)
|
|
bpy.types.Scene.comfyui_status = bpy.props.StringProperty(default="Ready")
|
|
bpy.types.Scene.comfyui_model_info = bpy.props.StringProperty(default="")
|
|
bpy.types.Scene.comfyui_log = bpy.props.StringProperty(default="")
|
|
|
|
def unregister_scene_props():
|
|
del bpy.types.Scene.comfyui_exp_warning
|
|
del bpy.types.Scene.comfyui_exp_log
|
|
del bpy.types.Scene.comfyui_server_url
|
|
del bpy.types.Scene.comfyui_workflow_path
|
|
del bpy.types.Scene.comfyui_input_image
|
|
del bpy.types.Scene.comfyui_preview_image
|
|
del bpy.types.Scene.comfyui_status
|
|
del bpy.types.Scene.comfyui_model_info
|
|
del bpy.types.Scene.comfyui_log
|
|
|
|
# ------------------------------
|
|
# Queue / Logging functions
|
|
# ------------------------------
|
|
def set_status(status, progress_text=""):
|
|
_status_queue.put((status, progress_text))
|
|
_log_queue.put(f"[STATUS] {status} | {progress_text}")
|
|
|
|
def set_model_info(model_info: dict):
|
|
_model_info_queue.put(model_info)
|
|
_log_queue.put(f"[MODEL INFO] {model_info}")
|
|
|
|
def log(message):
|
|
_log_queue.put(f"[LOG] {message}")
|
|
|
|
def process_queues():
|
|
scene = bpy.context.scene
|
|
try:
|
|
while not _status_queue.empty():
|
|
status, progress_text = _status_queue.get_nowait()
|
|
scene.comfyui_status = status
|
|
|
|
while not _model_info_queue.empty():
|
|
info = _model_info_queue.get_nowait()
|
|
scene.comfyui_model_info = json.dumps(info)
|
|
|
|
while not _log_queue.empty():
|
|
entry = _log_queue.get_nowait()
|
|
# Put newest entry at the top
|
|
current_logs = scene.comfyui_log.split('\n')
|
|
new_logs = [entry] + current_logs
|
|
# Keep only the last 100 entries to maintain performance
|
|
scene.comfyui_log = "\n".join(new_logs[:100])
|
|
|
|
except Exception as e:
|
|
print(f"[ComfyUI] Queue error: {e}")
|
|
return 0.5
|
|
|
|
def get_workflow_path_from_preferences():
|
|
"""Get workflow path from addon preferences"""
|
|
try:
|
|
prefs = bpy.context.preferences.addons[__name__].preferences
|
|
return prefs.comfyui_workflow_path
|
|
except:
|
|
return ""
|
|
|
|
def get_workflow_path():
|
|
"""Get workflow path from scene or preferences"""
|
|
scene = bpy.context.scene
|
|
workflow_path = scene.comfyui_workflow_path
|
|
if not workflow_path:
|
|
workflow_path = get_workflow_path_from_preferences()
|
|
return workflow_path
|
|
|
|
def start_queue_timer():
|
|
if not hasattr(bpy.app.timers, "_comfyui_timer_running"):
|
|
bpy.app.timers.register(process_queues, first_interval=0.5)
|
|
bpy.app.timers._comfyui_timer_running = True
|
|
|
|
# ------------------------------
|
|
# Threaded model generation
|
|
# ------------------------------
|
|
def generate_model_thread(server_url, workflow_path, input_image):
|
|
try:
|
|
scene = bpy.context.scene
|
|
scene.comfyui_log = ""
|
|
set_status("Loading workflow...", "")
|
|
# If workflow_path is empty, try to get it from preferences
|
|
if not workflow_path:
|
|
workflow_path = get_workflow_path()
|
|
|
|
if not workflow_path or not os.path.exists(workflow_path):
|
|
set_status("Error", "Workflow file not found")
|
|
log("Workflow file not found")
|
|
return
|
|
|
|
workflow_data = {}
|
|
with open(workflow_path, "r") as f:
|
|
workflow_data = json.load(f)
|
|
|
|
# if "204" in workflow_data and "inputs" in workflow_data["204"]:
|
|
# workflow_data["204"]["inputs"]["image"] = input_image
|
|
input_node_id = find_node_by_class(workflow_data, "LoadImage")
|
|
if input_node_id:
|
|
workflow_data[input_node_id]["inputs"]["image"] = input_image
|
|
log(f"Injected image into node {input_node_id} (LoadImage)")
|
|
else:
|
|
log("Error: Could not find a 'LoadImage' node in the workflow!")
|
|
return
|
|
|
|
wrapper = {"prompt": workflow_data, "client_id": "ai-creative-studio-node"}
|
|
set_status("Sending to ComfyUI...", "")
|
|
r = requests.post(f"{server_url}/prompt", json=wrapper)
|
|
if r.status_code != 200:
|
|
set_status(f"Error {r.status_code}", r.text[:100])
|
|
return
|
|
|
|
prompt_id = r.json()["prompt_id"]
|
|
set_status("Processing...", f"Prompt ID {prompt_id}")
|
|
log(f"Prompt sent: {prompt_id}")
|
|
|
|
model_info = wait_for_completion(server_url, prompt_id, workflow_data)
|
|
if model_info:
|
|
set_model_info(model_info)
|
|
set_status("Ready to import", model_info["filename"])
|
|
log(f"Model ready: {model_info['filename']}")
|
|
|
|
download_url = model_info["model_url"]
|
|
temp_file = os.path.join(bpy.app.tempdir, model_info['filename'])
|
|
r = requests.get(download_url)
|
|
if r.status_code == 200:
|
|
with open(temp_file, "wb") as f:
|
|
f.write(r.content)
|
|
bpy.app.timers.register(lambda: import_model_file(temp_file))
|
|
else:
|
|
set_status("Error", f"Download failed: {r.status_code}")
|
|
log(f"Download failed: {r.status_code}")
|
|
else:
|
|
set_status("Error", "Generation timed out")
|
|
log("Generation timed out")
|
|
except Exception as e:
|
|
set_status("Error", str(e))
|
|
log(f"Generation exception: {e}")
|
|
|
|
def wait_for_completion(server_url, prompt_id, workflow_data):
|
|
import time
|
|
import requests
|
|
total_nodes = len(workflow_data)
|
|
max_tries = 120
|
|
delay = 5
|
|
|
|
# 1. Get the prefix from the node you titled "FINAL_EXPORT"
|
|
target_node_id = find_node_by_title(workflow_data, "FINAL_EXPORT")
|
|
|
|
# Fallback to a default if the title is missing, or handle error
|
|
target_prefix = "textured"
|
|
if target_node_id and "inputs" in workflow_data[target_node_id]:
|
|
# Get the prefix (e.g., "3D/Hy3D_textured")
|
|
full_prefix = workflow_data[target_node_id]["inputs"].get("filename_prefix", "")
|
|
# We just need the last part to match the filename
|
|
target_prefix = full_prefix.split("/")[-1].split("\\")[-1]
|
|
|
|
log(f"Waiting for model with prefix matching: {target_prefix}")
|
|
|
|
for i in range(max_tries):
|
|
try:
|
|
set_status("Processing...", f"Checking completion")
|
|
r = requests.get(f"{server_url}/history/{prompt_id}")
|
|
|
|
if r.status_code != 200:
|
|
time.sleep(delay)
|
|
continue
|
|
|
|
hist = r.json()
|
|
if prompt_id in hist:
|
|
set_status("Finishing", "Download starting...")
|
|
prompt_output = hist[prompt_id].get("outputs", {})
|
|
|
|
# 2. Search all outputs for a match to our target_prefix
|
|
for node_id, node_data in prompt_output.items():
|
|
if "result" in node_data and node_data["result"]:
|
|
for item in node_data["result"]:
|
|
if isinstance(item, str) and target_prefix in item:
|
|
model_path = item
|
|
|
|
# Normalize path for Blender/Python
|
|
normalized_path = model_path.replace("\\\\", "/").replace("\\", "/")
|
|
|
|
if "/" in normalized_path:
|
|
subfolder, filename = normalized_path.rsplit("/", 1)
|
|
else:
|
|
subfolder, filename = "", normalized_path
|
|
|
|
return {
|
|
"filename": filename,
|
|
"subfolder": subfolder,
|
|
"type": "output",
|
|
"model_url": f"{server_url}/view?filename={filename}&subfolder={subfolder}&type=output"
|
|
}
|
|
|
|
time.sleep(delay)
|
|
except Exception as e:
|
|
log(f"Error during completion check: {e}")
|
|
break
|
|
|
|
return None
|
|
|
|
def import_model_file(filepath):
|
|
"""Bridge function for the timer to trigger the import operator."""
|
|
try:
|
|
# Check if file exists before trying to run the operator
|
|
if os.path.exists(filepath):
|
|
bpy.ops.wm.comfyui_import_model()
|
|
|
|
for window in bpy.context.window_manager.windows:
|
|
for area in window.screen.areas:
|
|
if area.type == 'VIEW_3D':
|
|
# Find the main viewport region (the 'WINDOW')
|
|
for region in area.regions:
|
|
if region.type == 'WINDOW':
|
|
with bpy.context.temp_override(window=window, area=area, region=region):
|
|
bpy.ops.view3d.view_selected()
|
|
return None
|
|
else:
|
|
log(f"Import failed: File {filepath} not found.")
|
|
except Exception as e:
|
|
print(f"Timer import error: {e}")
|
|
|
|
return None # Timers must return None to stop
|
|
|
|
# ------------------------------
|
|
# Operators
|
|
# ------------------------------
|
|
class ComfyUIGenerateOperator(bpy.types.Operator):
|
|
bl_idname = "wm.comfyui_generate_model"
|
|
bl_label = "Generate Model"
|
|
|
|
def execute(self, context):
|
|
start_queue_timer()
|
|
scene = context.scene
|
|
# Get workflow path from scene or preferences
|
|
workflow_path = get_workflow_path()
|
|
|
|
if not scene.comfyui_server_url or not workflow_path or not os.path.exists(workflow_path):
|
|
self.report({'ERROR'}, "Server URL or workflow invalid")
|
|
return {'CANCELLED'}
|
|
threading.Thread(target=generate_model_thread, args=(
|
|
scene.comfyui_server_url,
|
|
workflow_path,
|
|
scene.comfyui_input_image
|
|
), daemon=True).start()
|
|
return {'FINISHED'}
|
|
|
|
class ComfyUIImportOperator(bpy.types.Operator):
|
|
bl_idname = "wm.comfyui_import_model"
|
|
bl_label = "Import Model"
|
|
|
|
def execute(self, context):
|
|
scene = context.scene
|
|
try:
|
|
# Get model info from the scene property
|
|
model_info_json = scene.comfyui_model_info
|
|
if not model_info_json:
|
|
self.report({'ERROR'}, "No model information available")
|
|
return {'CANCELLED'}
|
|
|
|
model_info = json.loads(model_info_json)
|
|
|
|
# Get the filename and subfolder
|
|
filename = model_info.get("filename", "")
|
|
subfolder = model_info.get("subfolder", "")
|
|
|
|
if not filename:
|
|
self.report({'ERROR'}, "Invalid model information")
|
|
return {'CANCELLED'}
|
|
|
|
# Construct the full path
|
|
if subfolder:
|
|
full_path = os.path.join(subfolder, filename)
|
|
else:
|
|
full_path = filename
|
|
|
|
# Import the model using Blender's import functionality
|
|
# This assumes the model is in a supported format (glb, obj, gltf)
|
|
import_path = os.path.join(bpy.app.tempdir, filename)
|
|
|
|
# Check if the file exists in temp directory
|
|
if os.path.exists(import_path):
|
|
# Import the model based on file extension
|
|
if filename.lower().endswith(('.glb', '.gltf')):
|
|
try:
|
|
bpy.ops.import_scene.gltf(filepath=import_path)
|
|
except AttributeError:
|
|
# Fallback for older Blender versions
|
|
bpy.ops.import_scene.obj(filepath=import_path)
|
|
elif filename.lower().endswith('.obj'):
|
|
bpy.ops.import_scene.obj(filepath=import_path)
|
|
else:
|
|
self.report({'ERROR'}, f"Unsupported file format: {filename}")
|
|
return {'CANCELLED'}
|
|
|
|
set_status("Model imported successfully", filename)
|
|
log(f"Model imported: {filename}")
|
|
else:
|
|
self.report({'ERROR'}, f"Model file not found in temp directory: {filename}")
|
|
return {'CANCELLED'}
|
|
|
|
except Exception as e:
|
|
set_status("Error importing model", str(e))
|
|
log(f"Import error: {e}")
|
|
return {'CANCELLED'}
|
|
|
|
return {'FINISHED'}
|
|
|
|
class ComfyUIClearLogOperator(bpy.types.Operator):
|
|
bl_idname = "wm.comfyui_clear_log"
|
|
bl_label = "Clear Log"
|
|
|
|
def execute(self, context):
|
|
context.scene.comfyui_log = ""
|
|
return {'FINISHED'}
|
|
|
|
class ComfyUIPreferences(bpy.types.AddonPreferences):
|
|
bl_idname = __name__
|
|
|
|
comfyui_workflow_path: bpy.props.StringProperty(
|
|
name="Workflow File",
|
|
subtype='FILE_PATH',
|
|
default=""
|
|
)
|
|
|
|
def draw(self, context):
|
|
layout = self.layout
|
|
layout.prop(self, "comfyui_workflow_path")
|
|
|
|
def init_addon_preferences():
|
|
"""Initialize addon preferences with default values"""
|
|
try:
|
|
prefs = bpy.context.preferences.addons[__name__].preferences
|
|
if not prefs.comfyui_workflow_path:
|
|
prefs.comfyui_workflow_path = ""
|
|
except:
|
|
# If preferences are not available yet, skip initialization
|
|
pass
|
|
|
|
# ------------------------------
|
|
# Operator classes list
|
|
# ------------------------------
|
|
classes = [
|
|
ComfyUIGenerateOperator,
|
|
ComfyUIImportOperator,
|
|
ComfyUIClearLogOperator
|
|
]
|