468 lines
20 KiB
JavaScript
468 lines
20 KiB
JavaScript
document.addEventListener('DOMContentLoaded', () => {
|
|
const controlsContainer = document.getElementById('controls');
|
|
const generateButton = document.getElementById('generate-button');
|
|
const previewContainer = document.getElementById('preview');
|
|
const workflowDropdown = document.getElementById('workflow-dropdown');
|
|
const viewGalleryButton = document.getElementById('view-gallery-button'); // New
|
|
|
|
// Modal elements
|
|
const imageModal = document.getElementById('image-modal');
|
|
const modalImage = document.getElementById('modal-image');
|
|
const closeButton = document.getElementsByClassName('close-button')[0];
|
|
|
|
// Event listener for image preview click
|
|
previewContainer.addEventListener('click', () => {
|
|
const currentImage = previewContainer.querySelector('img');
|
|
if (currentImage) {
|
|
modalImage.src = currentImage.src;
|
|
imageModal.style.display = 'flex';
|
|
}
|
|
});
|
|
|
|
// Event listener for close button
|
|
closeButton.addEventListener('click', () => {
|
|
imageModal.style.display = 'none';
|
|
});
|
|
|
|
// Event listener for clicking outside the image to close the modal
|
|
imageModal.addEventListener('click', (event) => {
|
|
if (event.target === imageModal) {
|
|
imageModal.style.display = 'none';
|
|
}
|
|
});
|
|
|
|
let workflow = null;
|
|
const client_id = Math.random().toString(36).substring(2);
|
|
let socket = null;
|
|
|
|
function setupWebSocket() {
|
|
socket = new WebSocket(`ws://${window.location.host}/ws`);
|
|
|
|
socket.onopen = () => {
|
|
console.log('WebSocket connection established.');
|
|
};
|
|
|
|
socket.onmessage = (event) => {
|
|
const data = JSON.parse(event.data);
|
|
console.log("CozyGen: Raw WebSocket message received:", data); // Keep general log for now
|
|
|
|
if (data.type === 'cozygen_image_ready') { // Listen for our custom message
|
|
console.log("CozyGen: Custom image ready received:", data.data);
|
|
const imageUrl = data.data.image_url;
|
|
const img = document.createElement('img');
|
|
img.src = imageUrl;
|
|
previewContainer.innerHTML = '';
|
|
previewContainer.appendChild(img);
|
|
}
|
|
};
|
|
|
|
socket.onclose = () => {
|
|
console.log('WebSocket connection closed. Reconnecting...');
|
|
setTimeout(setupWebSocket, 1000);
|
|
};
|
|
}
|
|
|
|
// New function to populate the dropdown
|
|
async function populateWorkflowDropdown() {
|
|
try {
|
|
const response = await fetch('/cozygen/workflows');
|
|
if (!response.ok) {
|
|
throw new Error('Failed to load workflow list.');
|
|
}
|
|
const data = await response.json();
|
|
const workflows = data.workflows;
|
|
|
|
workflowDropdown.innerHTML = '';
|
|
|
|
if (workflows.length === 0) {
|
|
// Handle case with no workflows
|
|
const option = document.createElement('option');
|
|
option.textContent = 'No workflows found in /workflows folder';
|
|
option.disabled = true;
|
|
workflowDropdown.appendChild(option);
|
|
return null; // No default workflow
|
|
}
|
|
|
|
// Add all workflow options
|
|
workflows.forEach(wf => {
|
|
const option = document.createElement('option');
|
|
option.value = wf;
|
|
option.textContent = wf;
|
|
workflowDropdown.appendChild(option);
|
|
});
|
|
|
|
// Set initial selection to the first workflow
|
|
const defaultWorkflow = workflows[0];
|
|
workflowDropdown.value = defaultWorkflow;
|
|
return defaultWorkflow; // Return the name of the default workflow
|
|
} catch (error) {
|
|
console.error('Error populating workflow dropdown:', error);
|
|
workflowDropdown.innerHTML = '';
|
|
const option = document.createElement('option');
|
|
option.textContent = 'Error loading workflows';
|
|
option.disabled = true;
|
|
workflowDropdown.appendChild(option);
|
|
return null;
|
|
}
|
|
}
|
|
|
|
// Modified loadWorkflow to accept a filename
|
|
async function loadWorkflow(workflowFileName) {
|
|
if (!workflowFileName) {
|
|
controlsContainer.innerHTML = `<p>Please select a workflow.</p>`;
|
|
return;
|
|
}
|
|
try {
|
|
// ALWAYS fetch from the dedicated workflows endpoint for consistency.
|
|
const fetchUrl = `/cozygen/workflows/${workflowFileName}`;
|
|
|
|
const response = await fetch(fetchUrl);
|
|
if (!response.ok) {
|
|
throw new Error(`Failed to load workflow: ${workflowFileName}`);
|
|
}
|
|
workflow = await response.json();
|
|
renderControls();
|
|
} catch (error) {
|
|
console.error(error);
|
|
controlsContainer.innerHTML = `<p>Error loading workflow: ${workflowFileName}. Make sure it's a valid workflow JSON.</p>`;
|
|
}
|
|
}
|
|
|
|
function createParameterControl(id, node) {
|
|
const controlWrapper = document.createElement('div');
|
|
controlWrapper.className = 'control-wrapper';
|
|
|
|
const { title, type, default: defaultValue, min, max, choices: choicesString } = node.inputs;
|
|
|
|
const label = document.createElement('label');
|
|
label.textContent = title;
|
|
|
|
let input;
|
|
if (type === 'INT' || type === 'FLOAT') {
|
|
input = document.createElement('input');
|
|
input.type = 'number';
|
|
input.min = min;
|
|
input.max = max;
|
|
input.value = defaultValue;
|
|
if (type === 'FLOAT') input.step = '0.01';
|
|
} else if (type === 'DROPDOWN') {
|
|
input = document.createElement('select');
|
|
if (choicesString) {
|
|
choicesString.split(',').map(c => c.trim()).forEach(choice => {
|
|
const option = document.createElement('option');
|
|
option.value = choice;
|
|
option.textContent = choice;
|
|
input.appendChild(option);
|
|
});
|
|
}
|
|
input.value = defaultValue;
|
|
} else { // STRING
|
|
input = document.createElement('input');
|
|
input.type = 'text';
|
|
input.value = defaultValue;
|
|
}
|
|
|
|
input.dataset.nodeId = id;
|
|
controlWrapper.appendChild(label);
|
|
controlWrapper.appendChild(input);
|
|
|
|
return controlWrapper;
|
|
}
|
|
|
|
function createDynamicInputControl(id, node) {
|
|
const controlWrapper = document.createElement('div');
|
|
controlWrapper.className = 'control-wrapper';
|
|
|
|
const { param_name, param_type, default_value, min_value, max_value, choices, multiline, Multiline, add_randomize_toggle } = node.inputs;
|
|
|
|
const label = document.createElement('label');
|
|
label.textContent = param_name;
|
|
|
|
let input;
|
|
if (param_type === 'INT' || param_type === 'FLOAT') {
|
|
input = document.createElement('input');
|
|
input.type = 'number';
|
|
input.min = min_value;
|
|
input.max = max_value;
|
|
input.value = default_value;
|
|
if (param_type === 'FLOAT') input.step = '0.01';
|
|
} else if (param_type === 'DROPDOWN') {
|
|
input = document.createElement('select');
|
|
if (choices) {
|
|
choices.split(',').map(c => c.trim()).forEach(choice => {
|
|
const option = document.createElement('option');
|
|
option.value = choice;
|
|
option.textContent = choice;
|
|
input.appendChild(option);
|
|
});
|
|
}
|
|
input.value = default_value;
|
|
} else if (param_type === 'BOOLEAN') {
|
|
input = document.createElement('input');
|
|
input.type = 'checkbox';
|
|
input.checked = (default_value.toLowerCase() === 'true');
|
|
} else { // STRING
|
|
if (multiline || Multiline) {
|
|
input = document.createElement('textarea');
|
|
input.rows = 3;
|
|
input.addEventListener('input', () => {
|
|
input.style.height = 'auto';
|
|
input.style.height = (input.scrollHeight) + 'px';
|
|
});
|
|
} else {
|
|
input = document.createElement('input');
|
|
input.type = 'text';
|
|
}
|
|
input.value = default_value;
|
|
}
|
|
|
|
if (label && input) {
|
|
input.dataset.nodeId = id;
|
|
input.dataset.inputName = 'default_value';
|
|
controlWrapper.appendChild(label);
|
|
controlWrapper.appendChild(input);
|
|
}
|
|
|
|
if ((param_type === 'INT' || param_type === 'FLOAT') && add_randomize_toggle) {
|
|
const randomizeContainer = document.createElement('div');
|
|
randomizeContainer.className = 'randomize-container'; // Use a more generic class name
|
|
const randomizeCheckbox = document.createElement('input');
|
|
randomizeCheckbox.type = 'checkbox';
|
|
randomizeCheckbox.id = `randomize-control-${id}`; // More generic ID
|
|
randomizeCheckbox.dataset.nodeId = id;
|
|
randomizeCheckbox.checked = false;
|
|
const randomizeLabel = document.createElement('label');
|
|
randomizeLabel.htmlFor = `randomize-control-${id}`;
|
|
randomizeLabel.textContent = 'Randomize'; // Simpler label
|
|
randomizeContainer.appendChild(randomizeCheckbox);
|
|
randomizeContainer.appendChild(randomizeLabel);
|
|
controlWrapper.appendChild(randomizeContainer);
|
|
}
|
|
|
|
return controlWrapper;
|
|
}
|
|
|
|
function renderControls() {
|
|
if (!workflow) return;
|
|
controlsContainer.innerHTML = '';
|
|
|
|
// Get all nodes that should be rendered as controls
|
|
const parameterNodes = Object.entries(workflow).filter(([, node]) => node.class_type === 'CozyGenParameter');
|
|
const dynamicInputNodes = Object.entries(workflow).filter(([, node]) => node.class_type === 'CozyGenDynamicInput');
|
|
|
|
// Combine them into a single list
|
|
const allControlNodes = [...parameterNodes, ...dynamicInputNodes];
|
|
|
|
// Sort the combined list by priority. Both node types are expected to have a 'priority' input.
|
|
allControlNodes.sort(([, a], [, b]) => {
|
|
const priorityA = a.inputs.priority || 0;
|
|
const priorityB = b.inputs.priority || 0;
|
|
return priorityA - priorityB;
|
|
});
|
|
|
|
// Process the sorted list to render controls, grouping where necessary
|
|
let i = 0;
|
|
while (i < allControlNodes.length) {
|
|
const [id, node] = allControlNodes[i];
|
|
const classType = node.class_type;
|
|
const type = (classType === 'CozyGenParameter') ? node.inputs.type : node.inputs.param_type;
|
|
|
|
let group = [];
|
|
if (type === 'INT' || type === 'FLOAT') {
|
|
// Check for a consecutive group of number inputs
|
|
for (let j = i; j < allControlNodes.length; j++) {
|
|
const [, groupNode] = allControlNodes[j];
|
|
const groupClassType = groupNode.class_type;
|
|
const groupType = (groupClassType === 'CozyGenParameter') ? groupNode.inputs.type : groupNode.inputs.param_type;
|
|
|
|
if (['INT', 'FLOAT'].includes(groupType)) {
|
|
group.push(allControlNodes[j]);
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (group.length > 1) {
|
|
// Render a group of controls side-by-side
|
|
const groupContainer = document.createElement('div');
|
|
groupContainer.className = 'control-group';
|
|
group.forEach(([groupId, groupNode]) => {
|
|
const createFn = (groupNode.class_type === 'CozyGenParameter') ? createParameterControl : createDynamicInputControl;
|
|
groupContainer.appendChild(createFn(groupId, groupNode));
|
|
});
|
|
controlsContainer.appendChild(groupContainer);
|
|
i += group.length; // Move index past the group
|
|
} else {
|
|
// Render a single, full-width control
|
|
const createFn = (classType === 'CozyGenParameter') ? createParameterControl : createDynamicInputControl;
|
|
const controlElement = createFn(id, node);
|
|
controlElement.classList.add('full-width');
|
|
controlsContainer.appendChild(controlElement);
|
|
i++; // Move to the next item
|
|
}
|
|
}
|
|
}
|
|
|
|
async function queuePrompt() {
|
|
if (!workflow) return;
|
|
|
|
// Create a deep copy of the workflow to avoid modifying the original
|
|
const workflowToSend = JSON.parse(JSON.stringify(workflow));
|
|
|
|
// Collect current parameter values from UI
|
|
const parameterValues = {};
|
|
// Select all relevant input elements
|
|
const inputs = controlsContainer.querySelectorAll('input[type="text"], input[type="number"], input[type="checkbox"], select, textarea');
|
|
inputs.forEach(input => {
|
|
const nodeId = input.dataset.nodeId;
|
|
const inputName = input.dataset.inputName || 'default'; // For CozyGenParameter, it's 'default', for CozyGenDynamicInput, it's 'default_value'
|
|
const node = workflow[nodeId];
|
|
|
|
let nodeType;
|
|
if (node.class_type === 'CozyGenDynamicInput') {
|
|
nodeType = node.inputs.param_type;
|
|
} else {
|
|
nodeType = node.inputs.type;
|
|
}
|
|
|
|
let value;
|
|
if (input.type === 'checkbox') {
|
|
value = input.checked;
|
|
} else {
|
|
value = input.value;
|
|
}
|
|
|
|
// Special handling for randomization
|
|
if (node.class_type === 'CozyGenDynamicInput' && (node.inputs.param_type === 'INT' || node.inputs.param_type === 'FLOAT')) {
|
|
const randomizeCheckbox = document.getElementById(`randomize-control-${nodeId}`);
|
|
if (randomizeCheckbox && randomizeCheckbox.checked) {
|
|
value = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER);
|
|
// Update the input field with the new random value
|
|
const numberInput = document.querySelector(`input[data-node-id="${nodeId}"][data-input-name="default_value"]`);
|
|
if(numberInput) numberInput.value = value;
|
|
}
|
|
}
|
|
|
|
// Convert value to correct type based on node definition
|
|
if (nodeType === 'INT') {
|
|
value = parseInt(value);
|
|
}
|
|
else if (nodeType === 'FLOAT') {
|
|
value = parseFloat(value);
|
|
}
|
|
else if (nodeType === 'BOOLEAN') {
|
|
value = (value === true || value === "true");
|
|
}
|
|
// DROPDOWN and STRING values are already strings
|
|
|
|
// Store the value, keyed by nodeId and inputName
|
|
if (!parameterValues[nodeId]) {
|
|
parameterValues[nodeId] = {};
|
|
}
|
|
parameterValues[nodeId][inputName] = value;
|
|
});
|
|
|
|
// Iterate through the workflow to inject parameter values and remove CozyGenParameter nodes
|
|
const nodesToDelete = [];
|
|
for (const nodeId in workflowToSend) {
|
|
const node = workflowToSend[nodeId];
|
|
|
|
// If it's a CozyGenParameter node, mark it for a deletion
|
|
if (node.class_type === 'CozyGenParameter') {
|
|
nodesToDelete.push(nodeId);
|
|
continue; // Skip processing inputs for CozyGenParameter nodes
|
|
}
|
|
|
|
// If it's a CozyGenDynamicInput node, update its default_value input
|
|
if (node.class_type === 'CozyGenDynamicInput') {
|
|
if (parameterValues[nodeId] && parameterValues[nodeId].default_value !== undefined) {
|
|
node.inputs.default_value = parameterValues[nodeId].default_value;
|
|
}
|
|
// No need to delete CozyGenDynamicInput nodes, as they remain in the workflow
|
|
continue; // Skip processing inputs for CozyGenDynamicInput nodes here
|
|
}
|
|
|
|
// Iterate through inputs of other nodes to replace CozyGenParameter references
|
|
if (node.inputs) {
|
|
for (const inputName in node.inputs) {
|
|
const inputValue = node.inputs[inputName];
|
|
|
|
// Check if the input value is a reference to a CozyGenParameter or CozyGenDynamicInput node
|
|
if (Array.isArray(inputValue) && inputValue.length === 2) {
|
|
const referencedNodeId = inputValue[0];
|
|
const referencedNode = workflowToSend[referencedNodeId];
|
|
|
|
if (referencedNode && referencedNode.class_type === 'CozyGenParameter' && Object.prototype.hasOwnProperty.call(parameterValues, referencedNodeId)) {
|
|
// Replace the reference with the actual value from CozyGenParameter
|
|
node.inputs[inputName] = parameterValues[referencedNodeId].default;
|
|
} else if (referencedNode && referencedNode.class_type === 'CozyGenDynamicInput' && Object.prototype.hasOwnProperty.call(parameterValues, referencedNodeId)) {
|
|
// Replace the reference with the actual value from CozyGenDynamicInput
|
|
node.inputs[inputName] = parameterValues[referencedNodeId].default_value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Delete CozyGenParameter nodes from the workflowToSend
|
|
nodesToDelete.forEach(id => {
|
|
delete workflowToSend[id];
|
|
});
|
|
|
|
// Fix for CozyGenOutput node: rename 'image' input to 'images'
|
|
for (const nodeId in workflowToSend) {
|
|
const node = workflowToSend[nodeId];
|
|
if (node.class_type === 'CozyGenOutput' && node.inputs && node.inputs.image) {
|
|
node.inputs.images = node.inputs.image;
|
|
delete node.inputs.image;
|
|
}
|
|
}
|
|
|
|
const prompt = {
|
|
prompt: workflowToSend,
|
|
client_id: client_id
|
|
};
|
|
|
|
console.log('Prompt object being sent:', prompt); // Log the prompt object
|
|
|
|
try {
|
|
const response = await fetch('/prompt', {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json'
|
|
},
|
|
body: JSON.stringify(prompt)
|
|
});
|
|
|
|
if (!response.ok) {
|
|
throw new Error('Failed to queue prompt');
|
|
}
|
|
|
|
const data = await response.json();
|
|
console.log('Prompt queued:', data);
|
|
|
|
} catch (error) {
|
|
console.error(error);
|
|
}
|
|
}
|
|
|
|
generateButton.addEventListener('click', queuePrompt);
|
|
workflowDropdown.addEventListener('change', (event) => {
|
|
loadWorkflow(event.target.value);
|
|
});
|
|
viewGalleryButton.addEventListener('click', () => {
|
|
window.location.href = 'gallery.html';
|
|
});
|
|
|
|
async function initializeApp() {
|
|
const defaultWorkflow = await populateWorkflowDropdown();
|
|
if (defaultWorkflow) {
|
|
await loadWorkflow(defaultWorkflow);
|
|
}
|
|
setupWebSocket();
|
|
}
|
|
|
|
initializeApp();
|
|
});
|