# Automation Stage 6 - Image Generation Fix Plan

**Date:** December 4, 2025  
**Status:** Analysis Complete - Implementation Required  
**Priority:** HIGH

---

## 🔍 PROBLEM IDENTIFICATION

### Current Issue

Stage 6 of the automation pipeline (Image Prompts → Generated Images) is **NOT running correctly**. The issue stems from using the wrong AI function for image generation.

### Root Cause Analysis

**Current Implementation (INCORRECT):**
```python
# File: backend/igny8_core/business/automation/services/automation_service.py
# Line ~935

engine = AIEngine(account=self.account)
result = engine.execute(
    fn=GenerateImagesFunction(),
    payload={'image_ids': [image.id]}  # ❌ WRONG
)
```

**Why It Fails:**

1. `GenerateImagesFunction()` expects:
   - Input: `{'ids': [task_ids]}`  (Task IDs, NOT Image IDs)
   - Purpose: Extract prompts from Tasks and generate images for tasks
   - Use case: When you have Tasks with content but no images

2. Automation Stage 6 has:
   - Input: Images records with `status='pending'` (already have prompts)
   - Purpose: Generate actual image URLs from existing prompts
   - Context: Images were created in Stage 5 by `GenerateImagePromptsFunction`

### How Other Stages Work Correctly

**Stage 1:** Keywords → Clusters
```python
engine.execute(
    fn=AutoClusterFunction(),
    payload={'keyword_ids': keyword_ids}  # ✅ Correct
)
```

**Stage 2:** Clusters → Ideas
```python
engine.execute(
    fn=GenerateIdeasFunction(),
    payload={'cluster_ids': cluster_ids}  # ✅ Correct
)
```

**Stage 4:** Tasks → Content
```python
engine.execute(
    fn=GenerateContentFunction(),
    payload={'ids': task_ids}  # ✅ Correct
)
```

**Stage 5:** Content → Image Prompts
```python
engine.execute(
    fn=GenerateImagePromptsFunction(),
    payload={'ids': content_ids}  # ✅ Correct
)
```

**Stage 6:** Image Prompts → Images (BROKEN)
```python
# Currently uses GenerateImagesFunction (WRONG)
# Should use process_image_generation_queue (CORRECT)
```

---

## ✅ CORRECT SOLUTION

### The Right Approach

**Use `process_image_generation_queue` Celery task** - This is the same approach used by:
1. Writer/Images page (`/writer/images/generate_images/` endpoint)
2. Manual image generation from prompts

**Evidence from Working Code:**

```python
# File: backend/igny8_core/modules/writer/views.py
# ImagesViewSet.generate_images()

from igny8_core.ai.tasks import process_image_generation_queue

task = process_image_generation_queue.delay(
    image_ids=image_ids,      # ✅ Accepts image_ids
    account_id=account_id,
    content_id=content_id
)
```

**What `process_image_generation_queue` Does:**

1. ✅ Accepts `image_ids` (list of Image record IDs)
2. ✅ Each Image record already has a `prompt` field (populated by Stage 5)
3. ✅ Generates images sequentially with progress tracking
4. ✅ Updates Images records: `status='pending'` → `status='generated'`
5. ✅ Downloads and saves images locally
6. ✅ Automatically handles credits deduction
7. ✅ Supports multiple providers (OpenAI, Runware)
8. ✅ Handles errors gracefully (continues on failure)

---

## 📋 IMPLEMENTATION PLAN

### Changes Required

**File:** `backend/igny8_core/business/automation/services/automation_service.py`

**Location:** `run_stage_6()` method (lines 874-1022)

### Step 1: Import the Correct Task

**Current:**
```python
from igny8_core.ai.functions.generate_images import GenerateImagesFunction
```

**Add:**
```python
from igny8_core.ai.tasks import process_image_generation_queue
```

### Step 2: Modify Stage 6 Logic

**Replace this block (lines ~920-945):**

```python
# INCORRECT - Delete this
for idx, image in enumerate(image_list, 1):
    try:
        content_title = image.content.title if image.content else 'Unknown'
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Generating image {idx}/{total_images}: {image.image_type} for '{content_title}'"
        )
        
        # Call AI function via AIEngine
        engine = AIEngine(account=self.account)
        result = engine.execute(
            fn=GenerateImagesFunction(),
            payload={'image_ids': [image.id]}  # ❌ WRONG
        )
        
        # Monitor task
        task_id = result.get('task_id')
        if task_id:
            self._wait_for_task(task_id, stage_number, f"Image for '{content_title}'", continue_on_error=True)
        
        images_processed += 1
```

**With this:**

```python
# CORRECT - Use process_image_generation_queue
for idx, image in enumerate(image_list, 1):
    try:
        content_title = image.content.title if image.content else 'Unknown'
        self.logger.log_stage_progress(
            self.run.run_id, self.account.id, self.site.id,
            stage_number, f"Generating image {idx}/{total_images}: {image.image_type} for '{content_title}'"
        )
        
        # Call process_image_generation_queue directly (same as Writer/Images page)
        from igny8_core.ai.tasks import process_image_generation_queue
        
        # Queue the task
        if hasattr(process_image_generation_queue, 'delay'):
            task = process_image_generation_queue.delay(
                image_ids=[image.id],
                account_id=self.account.id,
                content_id=image.content.id if image.content else None
            )
            task_id = str(task.id)
        else:
            # Fallback for testing (synchronous)
            result = process_image_generation_queue(
                image_ids=[image.id],
                account_id=self.account.id,
                content_id=image.content.id if image.content else None
            )
            task_id = None
        
        # Monitor task (if async)
        if task_id:
            self._wait_for_task(task_id, stage_number, f"Image for '{content_title}'", continue_on_error=True)
        
        images_processed += 1
```

### Step 3: Update Logging

The logging structure remains the same, just update the log messages to reflect the correct process:

```python
self.logger.log_stage_progress(
    self.run.run_id, self.account.id, self.site.id,
    stage_number, f"Image generation task queued for '{content_title}' ({images_processed}/{total_images})"
)
```

### Step 4: No Changes Needed For

✅ Stage 5 (Image Prompt Extraction) - Already correct  
✅ Images table structure - Already has all required fields  
✅ Progress tracking - Already implemented in `process_image_generation_queue`  
✅ Credits deduction - Automatic in `process_image_generation_queue`  
✅ Error handling - Built into the task with `continue_on_error=True`

---

## 🔄 HOW IT WORKS (CORRECTED FLOW)

### Stage 5: Content → Image Prompts

```
Input:  Content (status='draft', no images)
AI:     GenerateImagePromptsFunction
Output: Images (status='pending', prompt='...')
```

### Stage 6: Image Prompts → Generated Images (FIXED)

```
Input:  Images (status='pending', has prompt)
Task:   process_image_generation_queue (Celery task)
AI:     Calls OpenAI/Runware API with prompt
Output: Images (status='generated', image_url='https://...', image_path='/path/to/file')
```

### What Happens in process_image_generation_queue

1. **Load Image Record:**
   - Get Image by ID
   - Read existing `prompt` field (created in Stage 5)
   - Get Content for template formatting

2. **Format Prompt:**
   - Use image_prompt_template from PromptRegistry
   - Format: `"Create a {image_type} image for '{post_title}'. Prompt: {image_prompt}"`
   - Handle model-specific limits (DALL-E 3: 4000 chars, DALL-E 2: 1000 chars)

3. **Generate Image:**
   - Call `AICore.generate_image()`
   - Uses configured provider (OpenAI/Runware)
   - Uses configured model (dall-e-3, runware:97@1, etc.)
   - Respects image size settings

4. **Download & Save:**
   - Download image from URL
   - Save to `/data/app/igny8/frontend/public/images/ai-images/`
   - Update Image record with both `image_url` and `image_path`

5. **Update Status:**
   - `status='pending'` → `status='generated'`
   - Triggers automatic Content status update (if all images generated)

6. **Deduct Credits:**
   - Automatic via `AICore` credit system
   - Records in `AIUsageLog`

---

## 🧪 TESTING CHECKLIST

### Pre-Deployment Tests

- [ ] **Unit Test:** Verify `process_image_generation_queue` works with single image
- [ ] **Integration Test:** Run Stage 6 with 3-5 pending images
- [ ] **Error Handling:** Test with invalid image ID
- [ ] **Credits:** Verify credits are deducted correctly
- [ ] **Multi-Provider:** Test with both OpenAI and Runware

### Post-Deployment Validation

- [ ] **Full Pipeline:** Run Automation from Stage 1 → Stage 7
- [ ] **Verify Stage 5 Output:** Images created with `status='pending'` and prompts
- [ ] **Verify Stage 6 Output:** Images updated to `status='generated'` with URLs
- [ ] **Check Downloads:** Images saved to `/data/app/igny8/frontend/public/images/ai-images/`
- [ ] **Monitor Logs:** Review automation logs for Stage 6 completion
- [ ] **Credits Report:** Confirm Stage 6 credits recorded in automation results

### Success Criteria

✅ Stage 6 completes without errors  
✅ All pending images get generated  
✅ Images are downloaded and accessible  
✅ Content status automatically updates when all images generated  
✅ Credits are properly deducted and logged  
✅ Automation proceeds to Stage 7 (Manual Review Gate)

---

## 📊 COMPARISON: BEFORE vs AFTER

### BEFORE (Broken)

```python
# ❌ WRONG APPROACH
GenerateImagesFunction()
- Expects: task_ids
- Purpose: Extract prompts from Tasks
- Problem: Doesn't work with Images that already have prompts
```

**Result:** Stage 6 fails, images never generated

### AFTER (Fixed)

```python
# ✅ CORRECT APPROACH
process_image_generation_queue()
- Accepts: image_ids
- Purpose: Generate images from existing prompts
- Works with: Images (status='pending' with prompts)
```

**Result:** Stage 6 succeeds, images generated sequentially with progress tracking

---

## 🔒 SAFETY & ROLLBACK

### Backup Plan

If the fix causes issues:

1. **Rollback Code:**
   - Git revert the automation_service.py changes
   - Automation still works for Stages 1-5

2. **Manual Workaround:**
   - Users can manually generate images from Writer/Images page
   - This uses the same `process_image_generation_queue` task

3. **No Data Loss:**
   - Stage 5 already created Images with prompts
   - These remain in database and can be processed anytime

---

## 📝 IMPLEMENTATION STEPS

1. **Update Code:** Modify `run_stage_6()` as documented above
2. **Test Locally:** Run automation with test data
3. **Code Review:** Verify changes match working Writer/Images implementation
4. **Deploy:** Push to production
5. **Monitor:** Watch first automation run for Stage 6 completion
6. **Validate:** Check images generated and credits deducted

---

## 🎯 EXPECTED OUTCOME

After implementing this fix:

✅ **Stage 6 will work correctly** - Images generate from prompts  
✅ **Consistent with manual flow** - Same logic as Writer/Images page  
✅ **Proper credits tracking** - Automated deduction via AICore  
✅ **Sequential processing** - One image at a time with progress  
✅ **Error resilience** - Continues on failure, logs errors  
✅ **Full pipeline completion** - Automation flows from Stage 1 → Stage 7  

---

## 🔗 RELATED FUNCTIONS

### Keep These Functions (Working Correctly)

- `GenerateImagePromptsFunction` - Stage 5 ✅
- `AutoClusterFunction` - Stage 1 ✅
- `GenerateIdeasFunction` - Stage 2 ✅
- `GenerateContentFunction` - Stage 4 ✅

### Use This Task for Stage 6

- `process_image_generation_queue` - Celery task for Images → Generated Images ✅

### DO NOT USE in Automation

- `GenerateImagesFunction` - For Tasks, not for Images with existing prompts ❌

---

## END OF PLAN

This plan provides a clear, actionable fix for Automation Stage 6 image generation, aligning it with the working manual image generation flow used throughout the application.