stable-diffusion.cpp-rest/test_inpainting_endpoint.cpp

#include <iostream>
#include <string>
#include <thread>
#include <chrono>
#include <nlohmann/json.hpp>
#include <httplib.h>
#include "include/utils.h"

using json = nlohmann::json;

// Simple test for the inpainting endpoint
int main() {
    const std::string serverUrl = "http://localhost:8080";
    httplib::Client client(serverUrl.c_str());

    std::cout << "Testing inpainting endpoint..." << std::endl;

    // Create a simple 1x1 white PNG (base64 encoded)
    // This is a minimal PNG header for a 1x1 white pixel
    const std::string whitePixelBase64 =
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8"
        "/5/hHgAHggJ/PchI7wAAAABJRU5ErkJggg==";

    // Create a simple 1x1 black PNG (base64 encoded) for mask
    const std::string blackPixelBase64 =
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk"
        "Y+h8DwADwA/wAA9AAB7BDQAAAABJRU5ErkJggg==";

    // Test inpainting request
    json inpaintingRequest = {
        {"prompt", "a red apple"},
        {"negative_prompt", "blurry, low quality"},
        {"source_image", whitePixelBase64},
        {"mask_image", blackPixelBase64},
        {"width", 512},
        {"height", 512},
        {"steps", 5}, // Use fewer steps for faster testing
        {"cfg_scale", 7.5},
        {"seed", "42"},
        {"sampling_method", "euler_a"},
        {"strength", 0.8}
    };

    std::cout << "Sending inpainting request..." << std::endl;

    auto res = client.Post("/api/generate/inpainting", inpaintingRequest.dump(), "application/json");

    if (res && res->status == 202) {
        std::cout << "✓ Inpainting request accepted (status 202)" << std::endl;

        try {
            json response = json::parse(res->body);
            std::string requestId = response.value("request_id", "");

            if (!requestId.empty()) {
                std::cout << "✓ Request ID: " << requestId << std::endl;

                // Poll for job completion
                std::cout << "Polling for job completion..." << std::endl;

                for (int i = 0; i < 60; ++i) { // Poll for up to 60 seconds
                    std::this_thread::sleep_for(std::chrono::seconds(1));

                    auto statusRes = client.Get(("/api/queue/job/" + requestId).c_str());

                    if (statusRes && statusRes->status == 200) {
                        json statusResponse = json::parse(statusRes->body);
                        std::string status = statusResponse["job"]["status"];

                        std::cout << "Job status: " << status << std::endl;

                        if (status == "completed") {
                            std::cout << "✓ Inpainting job completed successfully!" << std::endl;

                            if (statusResponse["job"].contains("outputs") &&
                                statusResponse["job"]["outputs"].is_array() &&
                                statusResponse["job"]["outputs"].size() > 0) {
                                std::cout << "✓ Generated " << statusResponse["job"]["outputs"].size()
                                          << " image(s)" << std::endl;
                            }

                            return 0; // Success
                        } else if (status == "failed") {
                            std::string error = statusResponse["job"].value("error_message", "Unknown error");
                            std::cerr << "✗ Inpainting job failed: " << error << std::endl;
                            return 1;
                        }
                    }
                }

                std::cerr << "✗ Job polling timeout" << std::endl;
                return 1;

            } else {
                std::cerr << "✗ No request ID in response" << std::endl;
                return 1;
            }
        } catch (const std::exception& e) {
            std::cerr << "✗ Failed to parse response: " << e.what() << std::endl;
            return 1;
        }
    } else {
        std::cerr << "✗ Failed to send inpainting request" << std::endl;
        if (res) {
            std::cerr << "Status: " << res->status << std::endl;
            std::cerr << "Response: " << res->body << std::endl;
        } else {
            std::cerr << "No response received" << std::endl;
        }
        return 1;
    }
}