custom_nodes/ComfyUI-FramerComfy/js/save-as-script.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { $el } from "../../scripts/ui.js";
import { framerComfyDialog } from "./input-dialog.js";



function getLoaderNodeKeys(workflow) {
	const loaderNodeTypes = [
		"CheckpointLoader",
		"CheckpointLoaderSimple",
		"DiffusersLoader",
		"unCLIPCheckpointLoader",
		"LoraLoader",
		"VAELoader",
		"ControlNetLoader",
		"DiffControlNetLoader",
		"UNETLoader",
		"CLIPLoader",
		"DualCLIPLoader",
		"CLIPVisionLoader",
		"StyleModelLoader",
		"GLIGENLoader",
		"UpscaleModelLoader",
		"HypernetworkLoader"
	];

	return Object.entries(workflow).filter(([key, node]) => loaderNodeTypes.includes(node.class_type))
}
function getFilenameFromHuggingFaceLink(url) {
	const urlParts = url.split('/');
	const repoIndex = urlParts.indexOf('huggingface.co') + 1;
	const repoIdParts = urlParts.slice(repoIndex, urlParts.indexOf('resolve'));
	const repo_id = repoIdParts.join('/');
	const file_name = urlParts[urlParts.length - 1];
	return { repo_id, file_name };
}



app.registerExtension({
	name: "Comfy.SaveAsScript",
	init() {
		$el("style", {
			parent: document.head,
		});
	},
	async setup() {
		async function savePythonScript() {
			// Use hardcoded API key for now
			var framer_comfy_api_key = "04bc0ebd-f8e5-48b5-92f6-72bb829ff76c"

			const loaderNodeSchemas = {
				"CheckpointLoader": {
					"ckpt_name": "Enter checkpoint url on Huggingface:",
				},
				"CheckpointLoaderSimple": {
					"ckpt_name": "Enter checkpoint url on Huggingface:",
				},
				"DiffusersLoader": {
					"model_path": "Enter model url on Huggingface:",
				},
				"unCLIPCheckpointLoader": {
					"ckpt_name": "Enter checkpoint url on Huggingface:",
				},
				"LoraLoader": {
					"lora_name": "Enter lora url on Huggingface:",
				},
				"VAELoader": {
					"vae_name": "Enter vae url on Huggingface:",
				},
				"ControlNetLoader": {
					"control_net_name": "Enter controlnet url on Huggingface:",
				},
				"DiffControlNetLoader": {
					"control_net_name": "Enter controlnet url on Huggingface:",
				},
				"UNETLoader": {
					"unet_name": "Enter unet url on Huggingface:",
				},
				"CLIPLoader": {
					"clip_name": "Enter clip url on Huggingface:",
				},
				"DualCLIPLoader": {
					"clip_name1": "Enter clip url on Huggingface:",
					"clip_name2": "Enter clip url on Huggingface:",
				},
				"CLIPVisionLoader": {
					"clip_name": "Enter clip url on Huggingface:",
				},
				"StyleModelLoader": {
					"style_model_name": "Enter model url on Huggingface:",
				},
				"GLIGENLoader": {
					"gligen_name": "Enter model url on Huggingface:",
				},
				"UpscaleModelLoader": {
					"model_name": "Enter model url on Huggingface:",
				},
				"HypernetworkLoader": {
					"hypernetwork_name": "Enter hypernetwork url on Huggingface:",
				}
			}


			const classtype_to_path = {
				"CheckpointLoader": "models/checkpoints",
				"CheckpointLoaderSimple": "models/checkpoints",
				"DiffusersLoader": "models/diffusers",
				"unCLIPCheckpointLoader": "models/checkpoints",
				"LoraLoader": "models/loras",
				"VAELoader": "models/vae",
				"ControlNetLoader": "models/controlnet",
				"DiffControlNetLoader": "models/controlnet",
				"UNETLoader": "models/unet",
				"CLIPLoader": "models/clip",
				"DualCLIPLoader": "models/clip",
				"CLIPVisionLoader": "models/clip_vision",
				"StyleModelLoader": "models/style_models",
				"GLIGENLoader": "models/gligen",
				"UpscaleModelLoader": "models/upscale_models",
				"HypernetworkLoader": "models/hypernetworks"
			}


			let workflow_models = []
			let workflow_name = "ComfyUI-FramerComfy"
			let huggingface_access_token = ""

			try {
				const p = await app.graphToPrompt();
				const loader_models = getLoaderNodeKeys(p.output);
				
				// Create input configuration for dialog
				const modelInputs = [];
				
				// Collect all required inputs for loader models
				for (const [nodeId, nodeData] of loader_models) {
					const nodeInputs = loaderNodeSchemas[nodeData.class_type];
					
					for (const [inputName, promptMessage] of Object.entries(nodeInputs)) {
						modelInputs.push({
							nodeId,
							nodeTitle: nodeData._meta.title,
							inputName,
							promptMessage
						});
					}
				}
				
				// Show the dialog even if we don't have models to get input for
				// because we need the workflow name and access token
				const dialogResponse = await framerComfyDialog.showModelUrlsDialog(modelInputs);
				
				// Process the dialog results
				if (dialogResponse && dialogResponse.workflowInfo) {
					// Get workflow information
					workflow_name = dialogResponse.workflowInfo.workflow_name;
					huggingface_access_token = dialogResponse.workflowInfo.huggingface_access_token;
					
					// Process model results
					for (const result of dialogResponse.modelResults) {
						const { nodeId, inputName, value } = result;
						const nodeData = p.output[nodeId];
						
						// Process the HuggingFace URL
						const { repo_id, file_name } = getFilenameFromHuggingFaceLink(value);
						const model_local_path = classtype_to_path[nodeData.class_type];
						
						// Update the node inputs and track the model
						p.output[nodeId].inputs[inputName] = file_name;
						workflow_models.push({ repo_id, file_name, model_local_path });
					}
				} else {
					// User canceled the dialog
					return;
				}
				
				// Create JSON payload
				const json = JSON.stringify({ 
					name: workflow_name,
					workflow_models: workflow_models, 
					framer_comfy_api_key: framer_comfy_api_key,
					hf_access_token: huggingface_access_token,
					workflow: JSON.stringify(p.output, null, 2) 
				}, null, 2);

				// Send to API
				var response = await api.fetchApi(`/saveasscript`, { method: "POST", body: json });
				if (response.status == 200) {
					const blob = new Blob([await response.text()], { type: "text/python;charset=utf-8" });
					const url = URL.createObjectURL(blob);
					
					let filename = workflow_name;
					if (!filename.endsWith(".py")) {
						filename += ".py";
					}

					const a = $el("a", {
						href: url,
						download: filename,
						style: { display: "none" },
						parent: document.body,
					});
					a.click();
					setTimeout(function () {
						a.remove();
						window.URL.revokeObjectURL(url);
					}, 0);
				}
			} catch (error) {
				console.error("Error in savePythonScript:", error);
				alert("Failed to save script: " + error.message);
			}
		}

		const menu = document.querySelector(".comfy-menu");
		const separator = document.createElement("hr");

		separator.style.margin = "20px 0";
		separator.style.width = "100%";
		menu.append(separator);

		const saveButton = document.createElement("button");
		saveButton.textContent = "Save as Script";
		saveButton.onclick = () => savePythonScript();
		menu.append(saveButton);


		// Also load to new style menu
		const dropdownMenu = document.querySelectorAll(".p-menubar-submenu ")[0];
		// Get submenu items
		const listItems = dropdownMenu.querySelectorAll("li");
		let newSetsize = listItems.length;

		const separatorMenu = document.createElement("li");
		separatorMenu.setAttribute("id", "pv_id_8_0_" + (newSetsize - 1).toString());
		separatorMenu.setAttribute("class", "p-menubar-separator");
		separatorMenu.setAttribute("role", "separator");
		separatorMenu.setAttribute("data-pc-section", "separator");

		dropdownMenu.append(separatorMenu);

		// Adjust list items within to increase setsize
		listItems.forEach((item) => {
			// First check if it's a separator
			if (item.getAttribute("data-pc-section") !== "separator") {
				item.setAttribute("aria-setsize", newSetsize);
			}
		});

		console.log(newSetsize);

		// Here's the format of list items		
		const saveButtonText = document.createElement("li");
		saveButtonText.setAttribute("id", "pv_id_8_0_" + newSetsize.toString());
		saveButtonText.setAttribute("class", "p-menubar-item relative");
		saveButtonText.setAttribute("role", "menuitem");
		saveButtonText.setAttribute("aria-label", "Deploy to FramerComfy");
		saveButtonText.setAttribute("aria-level", "2");
		saveButtonText.setAttribute("aria-setsize", newSetsize.toString());
		saveButtonText.setAttribute("aria-posinset", newSetsize.toString());
		saveButtonText.setAttribute("data-pc-section", "item");
		saveButtonText.setAttribute("data-p-active", "false");
		saveButtonText.setAttribute("data-p-focused", "false");

		saveButtonText.innerHTML = `
			<div class="p-menubar-item-content" data-pc-section="itemcontent">
				<a class="p-menubar-item-link" tabindex="-1" aria-hidden="true" data-pc-section="itemlink" target="_blank">
					<span class="p-menubar-item-icon pi pi-book"></span>
					<span class="p-menubar-item-label">Deploy to FramerComfy</span>
				</a>
			</div>
		`

		saveButtonText.onclick = () => savePythonScript();

		dropdownMenu.append(saveButtonText);



		console.log("SaveAsScript loaded");
	}
});