stable-diffusion.cpp-rest/include/generation_queue.h

#ifndef GENERATION_QUEUE_H
#define GENERATION_QUEUE_H

#include <string>
#include <memory>
#include <queue>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <future>
#include <thread>
#include <atomic>
#include <unordered_map>
#include <chrono>
#include <vector>
#include <map>

/**
 * @brief Job type enumeration
 */
enum class JobType {
    GENERATION,  ///< Image generation job
    HASHING,     ///< Model hashing job
    CONVERSION   ///< Model conversion/quantization job
};

/**
 * @brief Generation status enumeration
 */
enum class GenerationStatus {
    QUEUED,      ///< Request is queued and waiting to be processed
    MODEL_LOADING,  ///< Model is being loaded
    PROCESSING,     ///< Request is currently being processed
    COMPLETED,      ///< Request completed successfully
    FAILED          ///< Request failed during processing
};

/**
 * @brief Sampling method enumeration (matching stable-diffusion.cpp)
 */
enum class SamplingMethod {
    EULER,
    EULER_A,
    HEUN,
    DPM2,
    DPMPP2S_A,
    DPMPP2M,
    DPMPP2MV2,
    IPNDM,
    IPNDM_V,
    LCM,
    DDIM_TRAILING,
    TCD,
    DEFAULT  ///< Use model default
};

/**
 * @brief Scheduler enumeration (matching stable-diffusion.cpp)
 */
enum class Scheduler {
    DISCRETE,
    KARRAS,
    EXPONENTIAL,
    AYS,
    GITS,
    SMOOTHSTEP,
    SGM_UNIFORM,
    SIMPLE,
    DEFAULT  ///< Use model default
};

/**
 * @brief Generation request structure with all stable-diffusion.cpp parameters
 */
struct GenerationRequest {
    // Basic parameters
    std::string id;                    ///< Unique request ID
    std::string modelName;             ///< Name of the model to use
    std::string prompt;                ///< Text prompt for generation
    std::string negativePrompt;        ///< Negative prompt (optional)

    // Image parameters
    int width = 512;                  ///< Image width
    int height = 512;                 ///< Image height
    int batchCount = 1;               ///< Number of images to generate

    // Sampling parameters
    int steps = 20;                   ///< Number of diffusion steps
    float cfgScale = 7.5f;            ///< CFG scale
    SamplingMethod samplingMethod = SamplingMethod::DEFAULT;  ///< Sampling method
    Scheduler scheduler = Scheduler::DEFAULT;                 ///< Scheduler

    // Seed control
    std::string seed = "42";           ///< Seed for generation ("random" for random)

    // Model paths (for advanced usage)
    std::string clipLPath;            ///< Path to CLIP-L model
    std::string clipGPath;            ///< Path to CLIP-G model
    std::string clipVisionPath;        ///< Path to CLIP-Vision model
    std::string t5xxlPath;            ///< Path to T5-XXL model
    std::string qwen2vlPath;          ///< Path to Qwen2VL model
    std::string qwen2vlVisionPath;     ///< Path to Qwen2VL Vision model
    std::string diffusionModelPath;     ///< Path to standalone diffusion model
    std::string vaePath;               ///< Path to VAE model
    std::string taesdPath;             ///< Path to TAESD model
    std::string controlNetPath;        ///< Path to ControlNet model
    std::string embeddingDir;         ///< Path to embeddings directory
    std::string loraModelDir;         ///< Path to LoRA model directory

    // Advanced parameters
    int clipSkip = -1;                ///< CLIP skip layers
    std::vector<int> skipLayers = {7, 8, 9};  ///< Layers to skip for SLG
    float strength = 0.75f;            ///< Strength for img2img
    float controlStrength = 0.9f;      ///< ControlNet strength

    // Performance parameters
    int nThreads = -1;                ///< Number of threads (-1 for auto)
    bool offloadParamsToCpu = false;  ///< Offload parameters to CPU
    bool clipOnCpu = false;          ///< Keep CLIP on CPU
    bool vaeOnCpu = false;            ///< Keep VAE on CPU
    bool diffusionFlashAttn = false;   ///< Use flash attention
    bool diffusionConvDirect = false;  ///< Use direct convolution
    bool vaeConvDirect = false;       ///< Use direct VAE convolution

    // Output parameters
    std::string outputPath;            ///< Output path for generated images

    // Image-to-image parameters
    std::string initImagePath;         ///< Path to init image for img2img (can be file path or base64)
    std::vector<uint8_t> initImageData; ///< Init image data (decoded)
    int initImageWidth = 0;            ///< Init image width
    int initImageHeight = 0;           ///< Init image height
    int initImageChannels = 3;         ///< Init image channels

    // ControlNet parameters
    std::string controlImagePath;      ///< Path to control image for ControlNet
    std::vector<uint8_t> controlImageData; ///< Control image data (decoded)
    int controlImageWidth = 0;         ///< Control image width
    int controlImageHeight = 0;        ///< Control image height
    int controlImageChannels = 3;      ///< Control image channels

    // Upscaler parameters
    std::string esrganPath;            ///< Path to ESRGAN model for upscaling
    uint32_t upscaleFactor = 4;        ///< Upscale factor (2 or 4)

    // Inpainting parameters
    std::string maskImagePath;          ///< Path to mask image for inpainting
    std::vector<uint8_t> maskImageData; ///< Mask image data (decoded)
    int maskImageWidth = 0;             ///< Mask image width
    int maskImageHeight = 0;            ///< Mask image height
    int maskImageChannels = 1;          ///< Mask image channels (grayscale)

    // Request type
    enum class RequestType {
        TEXT2IMG,
        IMG2IMG,
        CONTROLNET,
        UPSCALER,
        INPAINTING
    } requestType = RequestType::TEXT2IMG;

    // Callback for completion
    std::function<void(const std::string&, const std::string&)> callback; ///< Callback for completion
};

/**
 * @brief Generation result structure
 */
struct GenerationResult {
    std::string requestId;             ///< ID of the original request
    GenerationStatus status;            ///< Final status of the generation
    bool success;                      ///< Whether generation was successful
    std::vector<std::string> imagePaths; ///< Paths to generated images (multiple for batch)
    std::string errorMessage;          ///< Error message if generation failed
    uint64_t generationTime;           ///< Time taken for generation in milliseconds
    int64_t actualSeed;               ///< Actual seed used for generation
};

/**
 * @brief Hash request structure for model hashing jobs
 */
struct HashRequest {
    std::string id;                    ///< Unique request ID
    std::vector<std::string> modelNames; ///< Model names to hash (empty = hash all unhashed)
    bool forceRehash = false;          ///< Force rehash even if hash exists
};

/**
 * @brief Hash result structure
 */
struct HashResult {
    std::string requestId;             ///< ID of the original request
    GenerationStatus status;            ///< Final status
    bool success;                      ///< Whether hashing was successful
    std::map<std::string, std::string> modelHashes; ///< Map of model names to their hashes
    std::string errorMessage;          ///< Error message if hashing failed
    uint64_t hashingTime;              ///< Time taken for hashing in milliseconds
    int modelsHashed;                  ///< Number of models successfully hashed
};

/**
 * @brief Conversion request structure for model quantization/conversion jobs
 */
struct ConversionRequest {
    std::string id;                    ///< Unique request ID
    std::string modelName;             ///< Model name to convert
    std::string modelPath;             ///< Full path to model file
    std::string outputPath;            ///< Output path for converted model
    std::string quantizationType;      ///< Quantization type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K)
};

/**
 * @brief Conversion result structure
 */
struct ConversionResult {
    std::string requestId;             ///< ID of the original request
    GenerationStatus status;            ///< Final status
    bool success;                      ///< Whether conversion was successful
    std::string outputPath;            ///< Path to converted model file
    std::string errorMessage;          ///< Error message if conversion failed
    uint64_t conversionTime;           ///< Time taken for conversion in milliseconds
    std::string originalSize;          ///< Original model file size
    std::string convertedSize;         ///< Converted model file size
};

/**
 * @brief Job information for queue status
 */
struct JobInfo {
    std::string id;                   ///< Job ID
    JobType type;                      ///< Job type (generation or hashing)
    GenerationStatus status;            ///< Current status
    std::string prompt;                ///< Job prompt (full text for generation, or model name for hashing)
    std::chrono::system_clock::time_point queuedTime;  ///< When job was queued
    std::chrono::system_clock::time_point startTime;   ///< When job started processing
    std::chrono::system_clock::time_point endTime;     ///< When job completed/failed
    int position;                      ///< Position in queue (for queued jobs)
    std::vector<std::string> outputFiles; ///< Paths to generated output files
    std::string errorMessage;          ///< Error message if job failed
    float progress = 0.0f;             ///< Overall progress (0.0 to 1.0)
    float modelLoadProgress = 0.0f;    ///< Model loading progress (0.0 to 1.0)
    float generationProgress = 0.0f;   ///< Generation progress (0.0 to 1.0)
    int currentStep = 0;               ///< Current step in generation
    int totalSteps = 0;                ///< Total steps in generation
    int64_t timeElapsed = 0;           ///< Time elapsed in milliseconds
    int64_t timeRemaining = 0;         ///< Estimated time remaining in milliseconds
    float speed = 0.0f;                ///< Generation speed in steps per second
    bool firstGenerationCallback = true; ///< Flag to track if this is the first generation callback

    // Enhanced fields for repeatable generation
    std::string modelName;             ///< Name of the model used
    std::string modelHash;             ///< SHA256 hash of the model
    std::string modelPath;             ///< Full path to the model file
    std::string negativePrompt;        ///< Negative prompt used
    int width = 512;                  ///< Image width
    int height = 512;                 ///< Image height
    int batchCount = 1;               ///< Number of images generated
    int steps = 20;                   ///< Number of diffusion steps
    float cfgScale = 7.5f;            ///< CFG scale
    SamplingMethod samplingMethod = SamplingMethod::DEFAULT;  ///< Sampling method used
    Scheduler scheduler = Scheduler::DEFAULT;                 ///< Scheduler used
    std::string seed;                 ///< Seed used for generation
    int64_t actualSeed = 0;           ///< Actual seed that was used (for random seeds)
    std::string requestType;          ///< Request type: text2img, img2img, controlnet, upscaler, inpainting
    float strength = 0.75f;           ///< Strength for img2img
    float controlStrength = 0.9f;     ///< ControlNet strength
    int clipSkip = -1;                ///< CLIP skip layers
    int nThreads = -1;                ///< Number of threads used
    bool offloadParamsToCpu = false;  ///< Offload parameters to CPU setting
    bool clipOnCpu = false;          ///< Keep CLIP on CPU setting
    bool vaeOnCpu = false;            ///< Keep VAE on CPU setting
    bool diffusionFlashAttn = false;   ///< Use flash attention setting
    bool diffusionConvDirect = false;  ///< Use direct convolution setting
    bool vaeConvDirect = false;       ///< Use direct VAE convolution setting
    uint64_t generationTime = 0;      ///< Total generation time in milliseconds

    // Image data for complex operations (base64 encoded)
    std::string initImageData;        ///< Init image data for img2img (base64)
    std::string controlImageData;     ///< Control image data for ControlNet (base64)
    std::string maskImageData;        ///< Mask image data for inpainting (base64)

    // Model paths for advanced usage
    std::string clipLPath;            ///< Path to CLIP-L model
    std::string clipGPath;            ///< Path to CLIP-G model
    std::string vaePath;              ///< Path to VAE model
    std::string taesdPath;            ///< Path to TAESD model
    std::string controlNetPath;       ///< Path to ControlNet model
    std::string embeddingDir;         ///< Path to embeddings directory
    std::string loraModelDir;         ///< Path to LoRA model directory
    std::string esrganPath;           ///< Path to ESRGAN model for upscaling
    uint32_t upscaleFactor = 4;       ///< Upscale factor used
};

/**
 * @brief Generation queue class for managing image generation requests
 *
 * This class manages a queue of image generation requests, processes them
 * asynchronously, and provides thread-safe access to the queue and results.
 * Only one generation job is processed at a time as specified in requirements.
 */
class GenerationQueue {
public:
    /**
     * @brief Construct a new Generation Queue object
     *
     * @param modelManager Pointer to the model manager
     * @param maxConcurrentGenerations Maximum number of concurrent generations (should be 1)
     * @param queueDir Directory to store job persistence files
     * @param outputDir Directory to store generated output files
     */
    explicit GenerationQueue(class ModelManager* modelManager, int maxConcurrentGenerations = 1,
                           const std::string& queueDir = "./queue", const std::string& outputDir = "./output");

    /**
     * @brief Destroy the Generation Queue object
     */
    virtual ~GenerationQueue();

    /**
     * @brief Add a generation request to the queue
     *
     * @param request The generation request
     * @return std::future<GenerationResult> Future for the generation result
     */
    std::future<GenerationResult> enqueueRequest(const GenerationRequest& request);

    /**
     * @brief Add a hash request to the queue
     *
     * @param request The hash request
     * @return std::future<HashResult> Future for the hash result
     */
    std::future<HashResult> enqueueHashRequest(const HashRequest& request);

    /**
     * @brief Add a conversion request to the queue
     *
     * @param request The conversion request
     * @return std::future<ConversionResult> Future for the conversion result
     */
    std::future<ConversionResult> enqueueConversionRequest(const ConversionRequest& request);

    /**
     * @brief Get the current queue size
     *
     * @return size_t Number of requests in the queue
     */
    size_t getQueueSize() const;

    /**
     * @brief Get the number of active generations
     *
     * @return size_t Number of currently processing requests
     */
    size_t getActiveGenerations() const;

    /**
     * @brief Get detailed queue status
     *
     * @return std::vector<JobInfo> List of all jobs with their status
     */
    std::vector<JobInfo> getQueueStatus() const;

    /**
     * @brief Get job information by ID
     *
     * @param jobId The job ID to look up
     * @return JobInfo Job information, or empty if not found
     */
    JobInfo getJobInfo(const std::string& jobId) const;

    /**
     * @brief Cancel a pending job
     *
     * @param jobId The job ID to cancel
     * @return true if job was cancelled, false if not found or already processing
     */
    bool cancelJob(const std::string& jobId);

    /**
     * @brief Clear all pending requests
     */
    void clearQueue();

    /**
     * @brief Start the queue processing thread
     */
    void start();

    /**
     * @brief Stop the queue processing thread
     */
    void stop();

    /**
     * @brief Check if the queue is running
     *
     * @return true if the queue is running, false otherwise
     */
    bool isRunning() const;

    /**
     * @brief Set the maximum number of concurrent generations
     *
     * @param maxConcurrent Maximum number of concurrent generations
     */
    void setMaxConcurrentGenerations(int maxConcurrent);

private:
    class Impl;
    std::unique_ptr<Impl> pImpl; // Pimpl idiom
};

#endif // GENERATION_QUEUE_H