diff --git a/Assets/runpod-backend.js b/Assets/runpod-backend.js
new file mode 100644
index 0000000..750e15e
--- /dev/null
+++ b/Assets/runpod-backend.js
@@ -0,0 +1,70 @@
+// RunPod Serverless Backend - Feature Handler
+// This file manages UI feature flags and parameter visibility for RunPod backend
+
+featureSetChangers.push(() => {
+ if (!gen_param_types) {
+ return [[], []];
+ }
+
+ const isRunPodBackend = currentModelHelper.curArch === 'runpod_serverless';
+
+ const coreParamsToHideForAPI = [
+ 'vaetilesize', 'vaetileoverlap', 'automaticvae',
+ 'clipstopatlayer', 'modelspecificenhancements'
+ ];
+
+ for (let param of gen_param_types) {
+ if (coreParamsToHideForAPI.includes(param.id)) {
+ if (isRunPodBackend) {
+ if (!param.hasOwnProperty('original_feature_flag_runpod')) {
+ param.original_feature_flag_runpod = param.feature_flag;
+ }
+ param.feature_flag = param.original_feature_flag_runpod
+ ? `${param.original_feature_flag_runpod},__runpod_incompatible__`
+ : '__runpod_incompatible__';
+ } else if (param.hasOwnProperty('original_feature_flag_runpod')) {
+ param.feature_flag = param.original_feature_flag_runpod;
+ delete param.original_feature_flag_runpod;
+ }
+ }
+ }
+
+ if (!isRunPodBackend) {
+ return [[], ['runpod_serverless']];
+ }
+
+ const removeFlags = [
+ 'sampling', 'refiners', 'controlnet', 'variation_seed',
+ 'video', 'autowebui', 'comfyui', 'frameinterps', 'ipadapter',
+ 'sdxl', 'cascade', 'sd3', 'seamless', 'freeu', 'teacache',
+ 'text2video', 'yolov8', 'aitemplate', 'endstepsearly',
+ 'dynamic_thresholding', 'zero_negative'
+ ];
+
+ // Features to add for RunPod backend
+ const addFlags = ['runpod_serverless', 'prompt', 'images'];
+
+ console.log(`[runpod-backend] Adding feature flags: ${addFlags.join(', ')}`);
+ console.log(`[runpod-backend] Removing feature flags: ${removeFlags.join(', ')}`);
+
+ return [addFlags, removeFlags];
+});
+
+if (typeof addModelChangeCallback === 'function') {
+ addModelChangeCallback(() => {
+ console.log(`[runpod-backend] Model changed to: ${currentModelHelper.curArch}`);
+
+ // Update the feature set and parameter visibility
+ reviseBackendFeatureSet();
+ hideUnsupportableParams();
+ });
+}
+
+// Initial parameter setup after UI loads
+setTimeout(() => {
+ console.log('[runpod-backend] Initial parameter setup starting');
+ reviseBackendFeatureSet();
+ hideUnsupportableParams();
+}, 500);
+
+console.log('[runpod-backend] Feature handler loaded');
diff --git a/RunPodServerlessBackend.cs b/RunPodServerlessBackend.cs
index a728b62..eb1a8d4 100644
--- a/RunPodServerlessBackend.cs
+++ b/RunPodServerlessBackend.cs
@@ -20,7 +20,10 @@ namespace Hartsy.Extensions.RunPodServerless;
public class RunPodServerlessBackend : AbstractT2IBackend
{
/// Cache of remote models by subtype.
- public ConcurrentDictionary> RemoteModels = null;
+ public ConcurrentDictionary> RemoteModels = new();
+
+ /// Model names by subtype (for internal tracking).
+ public new ConcurrentDictionary> Models = new();
public Session Session = null;
@@ -61,10 +64,19 @@ public class Settings : AutoConfiguration
[ConfigComment("Attempt to refresh models from worker on backend init")]
public bool AutoRefresh = false;
+
+ [ConfigComment("Worker SwarmUI port (default: 7801)")]
+ public int WorkerPort = 7801;
}
public Settings Config => (Settings)SettingsRaw;
+ /// Construct worker public URL from worker ID and port.
+ public string GetWorkerPublicUrl(string workerId)
+ {
+ return $"https://{workerId}-{Config.WorkerPort}.proxy.runpod.net";
+ }
+
/// Auto-throw exception if response indicates session error.
public static void AutoThrowException(JObject data)
{
@@ -394,8 +406,11 @@ public async Task RefreshModelsFromWorkerAsync(Session session = null)
modelMetadata[modelName] = new JObject
{
["name"] = modelName,
+ ["title"] = modelName,
["local"] = false,
- ["subtype"] = subtypeLocal
+ ["subtype"] = subtypeLocal,
+ ["architecture"] = "stable-diffusion-v1",
+ ["is_supported_model_format"] = true
};
}
}
@@ -417,8 +432,6 @@ public async Task RefreshModelsFromWorkerAsync(Session session = null)
int total = tempModels.Values.Sum(list => list.Count);
if (total > 0)
{
- RemoteModels ??= new ConcurrentDictionary>();
- Models ??= new ConcurrentDictionary>();
foreach (var kv in tempModels)
{
Models[kv.Key] = kv.Value;
@@ -428,6 +441,10 @@ public async Task RefreshModelsFromWorkerAsync(Session session = null)
RemoteModels[kv.Key] = kv.Value;
}
Logs.Debug($"[RunPodServerless] Model refresh complete: {total} models across {tempModels.Count} subtypes");
+ foreach (var kv in RemoteModels)
+ {
+ Logs.Verbose($"[RunPodServerless] RemoteModels['{kv.Key}'] = {kv.Value.Count} models");
+ }
return;
}
@@ -665,37 +682,72 @@ public async Task ClearWorkerStateAsync()
}
}
- /// Wake up worker and poll until ready.
+ /// Wake up worker and poll until ready. Uses job status to get worker ID, then probes worker URL directly.
public async Task WakeupAndWaitForWorkerAsync(RunPodApiClient client, int keepaliveDuration)
{
int keepaliveInterval = 30;
Logs.Verbose($"[RunPodServerless] Initiating wakeup with {keepaliveDuration}s keepalive...");
- await client.WakeupWorkerAsync(keepaliveDuration, keepaliveInterval);
- await Task.Delay(2000);
+ string jobId = await client.WakeupWorkerAsync(keepaliveDuration, keepaliveInterval);
int maxWaitSeconds = Config.StartupTimeoutSec;
int pollIntervalMs = Config.PollIntervalMs;
int maxAttempts = (maxWaitSeconds * 1000) / pollIntervalMs;
- Logs.Info($"[RunPodServerless] Polling for worker ready (max {maxWaitSeconds}s, interval {pollIntervalMs}ms)...");
+ string workerId = null;
+ string publicUrl = null;
+
+ AddLoadStatus("Waiting for RunPod worker to start...");
+
for (int attempt = 0; attempt < maxAttempts; attempt++)
{
try
{
- WorkerReadyResponse readyResponse = await client.CheckWorkerReadyAsync();
- if (readyResponse.Ready)
+ if (string.IsNullOrEmpty(workerId))
{
- int elapsedSeconds = (attempt * pollIntervalMs) / 1000;
- Logs.Info($"[RunPodServerless] Worker ready after {elapsedSeconds}s");
- return new WorkerInfo
+ JObject jobStatus = await client.GetJobStatusAsync(jobId);
+ string status = jobStatus["status"]?.ToString();
+ workerId = jobStatus["workerId"]?.ToString();
+
+ Logs.Verbose($"[RunPodServerless] Job {jobId} status: {status}, workerId: {workerId ?? "(none)"}");
+
+ if (status == "FAILED")
+ {
+ string error = jobStatus["error"]?.ToString() ?? "Job failed";
+ throw new Exception($"Wakeup job failed: {error}");
+ }
+
+ if (!string.IsNullOrEmpty(workerId))
{
- PublicUrl = readyResponse.PublicUrl,
- SessionId = readyResponse.SessionId,
- WorkerId = readyResponse.WorkerId,
- Version = readyResponse.Version
- };
+ publicUrl = GetWorkerPublicUrl(workerId);
+ Logs.Info($"[RunPodServerless] Worker ID obtained: {workerId}, URL: {publicUrl}");
+ AddLoadStatus($"Worker starting: {workerId}");
+ }
}
- if (!string.IsNullOrEmpty(readyResponse.Error))
+
+ if (!string.IsNullOrEmpty(publicUrl))
{
- Logs.Verbose($"[RunPodServerless] Worker not ready: {readyResponse.Error}");
+ try
+ {
+ JObject sessionReq = new() { ["local"] = true };
+ JObject sessionResp = await client.CallSwarmUIAsync(publicUrl, "/API/GetNewSession", sessionReq, timeoutSeconds: 10);
+
+ string sessionId = sessionResp["session_id"]?.ToString();
+ if (!string.IsNullOrEmpty(sessionId))
+ {
+ int elapsedSeconds = (attempt * pollIntervalMs) / 1000;
+ Logs.Info($"[RunPodServerless] Worker ready after {elapsedSeconds}s (session: {sessionId})");
+ AddLoadStatus("Worker ready!");
+ return new WorkerInfo
+ {
+ PublicUrl = publicUrl,
+ SessionId = sessionId,
+ WorkerId = workerId,
+ Version = sessionResp["version"]?.ToString()
+ };
+ }
+ }
+ catch (Exception probeEx)
+ {
+ Logs.Verbose($"[RunPodServerless] Worker probe failed (still starting): {probeEx.Message}");
+ }
}
}
catch (Exception ex)
diff --git a/RunPodServerlessExtension.cs b/RunPodServerlessExtension.cs
index 15b69b8..f15f68a 100644
--- a/RunPodServerlessExtension.cs
+++ b/RunPodServerlessExtension.cs
@@ -1,6 +1,7 @@
using System.Linq;
using System.Collections.Generic;
using FreneticUtilities.FreneticExtensions;
+using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Utils;
@@ -25,6 +26,8 @@ public class RunPodServerlessExtension : Extension
public override void OnPreInit()
{
Logs.Init("Initializing Hartsy's RunPod Serverless Backend Extension...");
+ // Register frontend script file for UI feature flags and parameter handling
+ ScriptFiles.Add("Assets/runpod-backend.js");
}
public override void OnInit()
@@ -59,13 +62,29 @@ public override void OnInit()
{
ModelsAPI.ExtraModelProviders["runpod_serverless"] = (string subtype) =>
{
- RunPodServerlessBackend[] backs = [.. Program.Backends.RunningBackendsOfType().Where(b => b.RemoteModels is not null)];
- IEnumerable> sets = backs.Select(b => b.RemoteModels.GetValueOrDefault(subtype)).Where(s => s is not null);
- if (!sets.Any())
+ Dictionary result = new();
+ var backends = Program.Backends.RunningBackendsOfType().ToList();
+ Logs.Verbose($"[RunPodServerless] ExtraModelProviders callback for subtype '{subtype}': found {backends.Count} running backend(s)");
+ foreach (var backend in backends)
{
- return [];
+ if (backend.RemoteModels != null && backend.RemoteModels.TryGetValue(subtype, out var models) && models != null)
+ {
+ Logs.Verbose($"[RunPodServerless] Backend has {models.Count} models for subtype '{subtype}'");
+ foreach (var kvp in models)
+ {
+ if (!result.ContainsKey(kvp.Key))
+ {
+ result[kvp.Key] = kvp.Value;
+ }
+ }
+ }
+ else
+ {
+ Logs.Verbose($"[RunPodServerless] Backend has no models for subtype '{subtype}' (RemoteModels null: {backend.RemoteModels == null})");
+ }
}
- return sets.Aggregate((a, b) => a.Union(b).PairsToDictionary(false));
+ Logs.Verbose($"[RunPodServerless] Returning {result.Count} total models for subtype '{subtype}'");
+ return result;
};
Logs.Debug("Registered RunPod Serverless models provider for extra remote models.");
}
diff --git a/WebAPI/RunPodApiClient.cs b/WebAPI/RunPodApiClient.cs
index 4130748..7569278 100644
--- a/WebAPI/RunPodApiClient.cs
+++ b/WebAPI/RunPodApiClient.cs
@@ -14,10 +14,11 @@ public class RunPodApiClient(string apiKey, string endpointId)
/// Shared HTTP client using SwarmUI's infrastructure.
public HttpClient HttpClient = NetworkBackendUtils.MakeHttpClient();
- /// Wake up worker with keepalive. Returns immediately after initiating wakeup.
+ /// Wake up worker with keepalive. Submits async job and returns job ID immediately.
/// How long to keep worker alive in seconds (default: 3600)
/// Ping interval in seconds (default: 30)
- public async Task WakeupWorkerAsync(int keepaliveDuration = 3600, int keepaliveInterval = 30, CancellationToken cancel = default)
+ /// Job ID for tracking the wakeup request
+ public async Task WakeupWorkerAsync(int keepaliveDuration = 3600, int keepaliveInterval = 30, CancellationToken cancel = default)
{
JObject payload = new()
{
@@ -29,7 +30,41 @@ public async Task WakeupWorkerAsync(int keepaliveDuration = 3600, int keepaliveI
}
};
Logs.Verbose($"[RunPodApiClient] Sending wakeup signal (duration: {keepaliveDuration}s, interval: {keepaliveInterval}s)");
- await CallRunPodHandlerAsync(payload, useSync: false, cancel);
+ string jobId = await SubmitJobAsync(payload, cancel);
+ Logs.Verbose($"[RunPodApiClient] Wakeup job submitted: {jobId}");
+ return jobId;
+ }
+
+ /// Submit a job to RunPod async endpoint, returns job ID immediately without waiting.
+ public async Task SubmitJobAsync(JObject payload, CancellationToken cancel = default)
+ {
+ string url = $"https://api.runpod.ai/v2/{endpointId}/run";
+ using HttpRequestMessage request = new(HttpMethod.Post, url)
+ {
+ Content = new StringContent(payload.ToString(), Encoding.UTF8, "application/json")
+ };
+ request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey);
+ Logs.Verbose($"[RunPodApiClient] Submitting async job to RunPod");
+ using HttpResponseMessage response = await HttpClient.SendAsync(request, cancel);
+ if (!response.IsSuccessStatusCode)
+ {
+ string error = await response.Content.ReadAsStringAsync(cancel);
+ throw new HttpRequestException($"RunPod API call failed ({response.StatusCode}): {error}");
+ }
+ string content = await response.Content.ReadAsStringAsync(cancel);
+ JObject result = JObject.Parse(content);
+ return result["id"]?.ToString() ?? throw new Exception("No job ID in RunPod response");
+ }
+
+ /// Get the current status of a job. Returns the full status response.
+ public async Task GetJobStatusAsync(string jobId, CancellationToken cancel = default)
+ {
+ string url = $"https://api.runpod.ai/v2/{endpointId}/status/{jobId}";
+ using HttpRequestMessage request = new(HttpMethod.Get, url);
+ request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey);
+ using HttpResponseMessage response = await HttpClient.SendAsync(request, cancel);
+ string content = await response.Content.ReadAsStringAsync(cancel);
+ return JObject.Parse(content);
}
/// Check if worker is ready for generation.
@@ -172,11 +207,13 @@ public async Task CallRunPodHandlerAsync(JObject payload, bool useSync,
return result["output"] as JObject ?? result;
}
- /// Poll job status for async RunPod calls.
+ /// Poll job status for async RunPod calls. Extracts worker info when IN_PROGRESS.
public async Task PollJobStatusAsync(string jobId, CancellationToken cancel)
{
string url = $"https://api.runpod.ai/v2/{endpointId}/status/{jobId}";
int maxAttempts = 300;
+ string lastWorkerId = null;
+
for (int attempt = 0; attempt < maxAttempts; attempt++)
{
cancel.ThrowIfCancellationRequested();
@@ -186,10 +223,27 @@ public async Task PollJobStatusAsync(string jobId, CancellationToken ca
string content = await response.Content.ReadAsStringAsync(cancel);
JObject result = JObject.Parse(content);
string status = result["status"]?.ToString();
+ string workerId = result["workerId"]?.ToString();
+
+ if (!string.IsNullOrEmpty(workerId) && workerId != lastWorkerId)
+ {
+ lastWorkerId = workerId;
+ Logs.Verbose($"[RunPodApiClient] Job {jobId} assigned to worker: {workerId}");
+ }
+
if (status == "COMPLETED")
{
+ // Ensure workerId is in the result for caller to use
+ if (!string.IsNullOrEmpty(lastWorkerId) && result["workerId"] == null)
+ {
+ result["workerId"] = lastWorkerId;
+ }
return result;
}
+ else if (status == "IN_PROGRESS")
+ {
+ Logs.Verbose($"[RunPodApiClient] Job {jobId} in progress on worker {workerId}...");
+ }
else if (status == "FAILED")
{
string error = result["error"]?.ToString() ?? "Job failed";