Free draws on Krea 2; VIP draws Krea 2 + MOSS-Image

crates/openplotva-app/src/image_jobs.rs

@@ -101,8 +101,9 @@ pub const IMAGE_REGULAR_JOB_WORKER_QUEUES: [&str; 1] = [IMAGE_REGULAR_QUEUE_NAME
 pub const IMAGE_EDIT_WORKFLOW_KEY: &str = "image_edit";
 pub const IMAGE_GENERATION_WORKFLOW_KEY: &str = "image_generation";
 pub const IMAGE_GENERATION_FLUX_WORKFLOW_KEY: &str = "image_generation_flux";
+/// Free users' text-to-image workflow; the key predates its current model.
 pub const IMAGE_GENERATION_BOOGU_TURBO_WORKFLOW_KEY: &str = "image_generation_boogu_turbo";
-/// Third VIP text-to-image slot, routed to MOSS-Image.
+/// Second VIP text-to-image slot, routed to MOSS-Image.
 pub const IMAGE_GENERATION_MOSS_WORKFLOW_KEY: &str = "image_generation_moss";
 pub const IMAGE_EDIT_FLUX_WORKFLOW_KEY: &str = "image_edit_flux";
 pub const IMAGE_EDIT_BOOGU_TURBO_WORKFLOW_KEY: &str = "image_edit_boogu_turbo";

crates/openplotva-app/src/lib.rs

@@ -149,6 +149,10 @@ const DURABLE_UPDATE_CONSUMER_MAX_WORKERS: usize =
     openplotva_storage::POSTGRES_CRITICAL_MAX_CONNECTIONS as usize * 3;
 const REGULAR_IMAGE_GENERATION_WORKFLOW_KEY: &str =
     image_jobs::IMAGE_GENERATION_BOOGU_TURBO_WORKFLOW_KEY;
+const VIP_IMAGE_GENERATION_WORKFLOW_KEYS: [&str; 2] = [
+    image_jobs::IMAGE_GENERATION_FLUX_WORKFLOW_KEY,
+    image_jobs::IMAGE_GENERATION_MOSS_WORKFLOW_KEY,
+];
 const VIP_IMAGE_EDITOR_WORKFLOW_KEY: &str = image_jobs::IMAGE_EDIT_FLUX_WORKFLOW_KEY;
 
 fn durable_update_consumer_worker_limit(configured: usize) -> usize {
@@ -13314,28 +13318,19 @@ async fn start_runtime_workers(
     )
     .with_openrouter_free_gate(Arc::clone(&openrouter_free_gate))
     .with_reporter(routing_event_reporter.clone());
-    let routed_vip_flux_image_generator = image_jobs::RoutedImageGenerator::new(
-        image_attempt_walker.clone(),
-        image_jobs::aifarm_draw_api_config_from_app_config(config),
-    )
-    .with_workflow_key(image_jobs::IMAGE_GENERATION_FLUX_WORKFLOW_KEY);
-    let routed_vip_boogu_image_generator = image_jobs::RoutedImageGenerator::new(
-        image_attempt_walker.clone(),
-        image_jobs::aifarm_draw_api_config_from_app_config(config),
-    )
-    .with_workflow_key(image_jobs::IMAGE_GENERATION_BOOGU_TURBO_WORKFLOW_KEY);
-    let routed_vip_moss_image_generator = image_jobs::RoutedImageGenerator::new(
-        image_attempt_walker.clone(),
-        image_jobs::aifarm_draw_api_config_from_app_config(config),
-    )
-    .with_workflow_key(image_jobs::IMAGE_GENERATION_MOSS_WORKFLOW_KEY);
+    let [vip_first_workflow_key, vip_second_workflow_key] = VIP_IMAGE_GENERATION_WORKFLOW_KEYS;
     let vip_image_generator = image_jobs::OptimizingImageGenerator::new(
         image_jobs::ParallelImageGenerator::new(
-            image_jobs::ParallelImageGenerator::new(
-                routed_vip_flux_image_generator,
-                routed_vip_boogu_image_generator,
-            ),
-            routed_vip_moss_image_generator,
+            image_jobs::RoutedImageGenerator::new(
+                image_attempt_walker.clone(),
+                image_jobs::aifarm_draw_api_config_from_app_config(config),
+            )
+            .with_workflow_key(vip_first_workflow_key),
+            image_jobs::RoutedImageGenerator::new(
+                image_attempt_walker.clone(),
+                image_jobs::aifarm_draw_api_config_from_app_config(config),
+            )
+            .with_workflow_key(vip_second_workflow_key),
         ),
         media_prompt_optimizer.clone(),
     )
@@ -15166,11 +15161,18 @@ mod tests {
     }
 
     #[test]
-    fn image_worker_workflows_keep_free_generation_on_boogu_and_editor_on_flux() {
+    fn image_worker_workflows_keep_free_vip_and_editor_slots_apart() {
         assert_eq!(
             super::REGULAR_IMAGE_GENERATION_WORKFLOW_KEY,
             super::image_jobs::IMAGE_GENERATION_BOOGU_TURBO_WORKFLOW_KEY
         );
+        assert_eq!(
+            super::VIP_IMAGE_GENERATION_WORKFLOW_KEYS,
+            [
+                super::image_jobs::IMAGE_GENERATION_FLUX_WORKFLOW_KEY,
+                super::image_jobs::IMAGE_GENERATION_MOSS_WORKFLOW_KEY,
+            ]
+        );
         assert_eq!(
             super::VIP_IMAGE_EDITOR_WORKFLOW_KEY,
             super::image_jobs::IMAGE_EDIT_FLUX_WORKFLOW_KEY

migrations/210_krea2_free_draws.down.sql

@@ -0,0 +1,9 @@
+LOCK TABLE workflow_assignments IN SHARE ROW EXCLUSIVE MODE;
+
+DELETE FROM workflow_assignments
+WHERE inference_overrides ? 'created_by_krea2_free_draws_v1';
+
+UPDATE workflow_assignments
+SET provider_model_id = (inference_overrides ->> 'krea2_free_draws_v1_previous_model_id')::BIGINT,
+    inference_overrides = inference_overrides - 'krea2_free_draws_v1_previous_model_id'
+WHERE inference_overrides ? 'krea2_free_draws_v1_previous_model_id';

migrations/210_krea2_free_draws.up.sql

@@ -0,0 +1,80 @@
+-- Free users' draws (`image_generation_boogu_turbo`) move from FLUX.3 to Krea 2 Turbo, with
+-- FLUX.3 as their fallback, as it is for the VIP Krea 2 slot. VIP draws drop their FLUX.3
+-- slot in code and keep Krea 2 and MOSS-Image; FLUX.3 keeps the VIP edits. Only a FLUX.3
+-- primary is switched, and only when an enabled Krea 2 model exists on aifarm-draw; any other
+-- database is left unchanged. The switched primary records its previous model and the created
+-- fallback carries a marker, so the down migration restores the previous routing exactly.
+
+LOCK TABLE llm_providers, provider_models, workflow_assignments
+    IN SHARE ROW EXCLUSIVE MODE;
+
+DO $$
+BEGIN
+    IF EXISTS (
+        SELECT 1 FROM workflow_assignments
+        WHERE inference_overrides ? 'krea2_free_draws_v1_previous_model_id'
+           OR inference_overrides ? 'created_by_krea2_free_draws_v1'
+    ) THEN
+        RAISE EXCEPTION 'Krea 2 free draws routing marker collision';
+    END IF;
+END $$;
+
+CREATE TEMPORARY TABLE krea2_free_draw_models ON COMMIT DROP AS
+SELECT
+    max(model.id) FILTER (WHERE model.config ->> 'endpoint_name' = 'krea2_generate') AS krea2_id,
+    max(model.id) FILTER (WHERE model.config ->> 'endpoint_name' = 'flux3_generate') AS flux3_id
+FROM provider_models AS model
+JOIN llm_providers AS provider ON provider.id = model.provider_id
+WHERE provider.name = 'aifarm-draw'
+  AND model.enabled;
+
+UPDATE workflow_assignments AS assignment
+SET provider_model_id = models.krea2_id,
+    inference_overrides = assignment.inference_overrides || jsonb_build_object(
+        'krea2_free_draws_v1_previous_model_id', assignment.provider_model_id
+    )
+FROM krea2_free_draw_models AS models
+WHERE assignment.workflow_key = 'image_generation_boogu_turbo'
+  AND assignment.scope = 'global'
+  AND assignment.role = 'primary'
+  AND assignment.enabled
+  AND assignment.provider_model_id = models.flux3_id
+  AND models.krea2_id IS NOT NULL;
+
+INSERT INTO workflow_assignments (
+    workflow_key,
+    scope,
+    role,
+    provider_model_id,
+    weight,
+    fallback_order,
+    enabled,
+    inference_overrides,
+    cb_failure_threshold,
+    cb_cooldown_ms
+)
+SELECT
+    'image_generation_boogu_turbo',
+    'global',
+    'fallback',
+    models.flux3_id,
+    NULL,
+    0,
+    TRUE,
+    '{"created_by_krea2_free_draws_v1": true}'::jsonb,
+    5,
+    30000
+FROM krea2_free_draw_models AS models
+WHERE EXISTS (
+      SELECT 1 FROM workflow_assignments
+      WHERE workflow_key = 'image_generation_boogu_turbo'
+        AND inference_overrides ? 'krea2_free_draws_v1_previous_model_id'
+  )
+  AND NOT EXISTS (
+      SELECT 1 FROM workflow_assignments
+      WHERE workflow_key = 'image_generation_boogu_turbo'
+        AND scope = 'global'
+        AND role = 'fallback'
+        AND enabled
+        AND provider_model_id = models.flux3_id
+  );