# Image Studio: Unified Architecture & Integration Patterns

**Purpose**: Define **reusable** code patterns and architecture for integrating 40+ WaveSpeed AI models into Image Studio  
**Status**: Architecture Proposal - Pre-Implementation Review  
**Based On**: Existing `main_image_generation.py` + Video Studio patterns  
**Key Principle**: **REUSABILITY** - Extend existing code, don't duplicate

---

## 📊 Executive Summary

This document proposes a **reusable architecture** for Image Studio that:
1. **✅ Extends Existing Code**: Builds on `main_image_generation.py` (already exists)
2. **✅ Extracts Reusable Helpers**: Validation and tracking from existing functions
3. **✅ Reuses Provider Pattern**: Extends `ImageGenerationProvider` protocol
4. **✅ Reuses Infrastructure**: WaveSpeedClient, validation, tracking logic
5. **✅ Scales to 40+ Models**: Easy addition by following existing patterns

---

## 🔍 Current State Analysis

### **Video Studio Pattern** (`main_video_generation.py`) - Reference

#### **Architecture**
```
┌─────────────────────────────────────────┐
│  ai_video_generate()                    │  ← Unified Entry Point
│  - Pre-flight validation                │
│  - Provider routing                     │
│  - Usage tracking                       │
│  - Progress callbacks                   │
└──────────────┬──────────────────────────┘
               │
       ┌───────┴────────┐
       │                │
┌──────▼──────┐  ┌─────▼──────────┐
│ HuggingFace │  │ WaveSpeed      │
│ Provider    │  │ Provider        │
└─────────────┘  └────────────────┘
```

#### **Key Patterns**
1. **Unified Entry Point**: `ai_video_generate()` handles all video operations
2. **Pre-flight Validation**: Subscription checks BEFORE API calls
3. **Provider Abstraction**: Routes to provider-specific handlers
4. **Standardized Returns**: Always returns `Dict[str, Any]` with consistent keys
5. **Usage Tracking**: Centralized `track_video_usage()` function
6. **Progress Callbacks**: Optional progress updates for async operations
7. **Error Handling**: Consistent HTTPException patterns

---

### **Image Studio Current Pattern** ✅ **ALREADY EXISTS**

#### **Architecture**
```
┌─────────────────────────────────────────┐
│  main_image_generation.py               │  ← Unified Entry Point (EXISTS)
│  - generate_image()                     │
│  - generate_character_image()           │
│  - Pre-flight validation                │
│  - Usage tracking                       │
└──────────────┬──────────────────────────┘
               │
    ┌──────────┼──────────┐
    │          │          │
┌───▼───┐ ┌───▼───┐ ┌───▼───┐
│Create │ │ Edit  │ │Upscale│
│Service│ │Service│ │Service│
└───┬───┘ └───┬───┘ └───┬───┘
    │          │          │
┌───▼──────────▼──────────▼───┐
│  image_generation/          │
│  - ImageGenerationProvider  │  ← Protocol (EXISTS)
│  - WaveSpeedImageProvider    │
│  - StabilityImageProvider    │
│  - HuggingFaceImageProvider  │
│  - GeminiImageProvider      │
└──────────────────────────────┘
```

#### **Current Implementation** ✅
1. **✅ Unified Entry Point EXISTS**: `main_image_generation.py` with `generate_image()`
2. **✅ Pre-flight Validation**: Implemented in `generate_image()`
3. **✅ Provider Abstraction**: `ImageGenerationProvider` protocol with implementations
4. **✅ Usage Tracking**: Implemented in `generate_image()`
5. **✅ Standardized Returns**: `ImageGenerationResult` dataclass

#### **Current Usage**
- ✅ **Used by**: YouTube, Podcast, Story Writer, Facebook Writer, LinkedIn
- ⚠️ **NOT used by**: `CreateStudioService` (uses providers directly)
- ⚠️ **Missing**: Editing, Upscaling, 3D operations don't use unified entry

#### **Reusability Opportunities**
1. **Extend `main_image_generation.py`** for editing operations
2. **Reuse provider pattern** for new WaveSpeed models
3. **Standardize all services** to use unified entry point
4. **Extract common validation/tracking** into reusable functions

---

## 🎯 Proposed Architecture Enhancement

### **Core Principle: Extend Existing Pattern for Maximum Reusability**

**Build on existing `main_image_generation.py`** instead of creating new modules. Extend it to support all image operations while maintaining the proven pattern.

### **Enhanced Architecture Diagram**

```
┌─────────────────────────────────────────────────────────────┐
│  main_image_generation.py (EXISTS - EXTEND)                 │
│  ✅ generate_image()              (text-to-image)           │
│  ✅ generate_character_image()    (character consistency)   │
│  🆕 generate_image_edit()          (editing operations)       │
│  🆕 generate_image_upscale()      (upscaling)               │
│  🆕 generate_image_to_3d()          (3D generation)           │
│  🆕 generate_face_swap()            (face swapping)            │
│  🆕 generate_image_translate()     (translation)              │
└──────────────┬──────────────────────────────────────────────┘
               │
    ┌──────────┼──────────┬──────────┐
    │          │          │          │
┌───▼───┐ ┌───▼───┐ ┌───▼───┐ ┌───▼───┐
│Generate│ │ Edit  │ │Upscale│ │Transform│
│Provider│ │Provider│ │Provider│ │Provider│
└───┬───┘ └───┬───┘ └───┬───┘ └───┬───┘
    │          │          │          │
┌───▼──────────▼──────────▼──────────▼───┐
│  image_generation/ (EXISTS - EXTEND)   │
│  ✅ ImageGenerationProvider Protocol   │
│  ✅ WaveSpeedImageProvider              │
│  🆕 WaveSpeedEditProvider               │
│  🆕 WaveSpeedUpscaleProvider             │
│  🆕 WaveSpeed3DProvider                  │
│  🆕 WaveSpeedFaceSwapProvider            │
└─────────────────────────────────────────┘
```

### **Key Reusability Principles**

1. **Reuse Existing Infrastructure**
   - Extend `main_image_generation.py` (don't duplicate)
   - Reuse `ImageGenerationProvider` protocol pattern
   - Reuse validation and tracking logic

2. **Consistent Function Signatures**
   - All functions follow same pattern: `generate_<operation>()`
   - All use same validation/tracking helpers
   - All return standardized results

3. **Provider Pattern Extension**
   - Create new provider classes following `ImageGenerationProvider` protocol
   - Reuse `WaveSpeedClient` for all WaveSpeed operations
   - Consistent error handling across providers

---

## 📐 Reusable Code Patterns

### **Pattern 1: Extend Existing Unified Entry Point** ✅

#### **Current Structure** (EXISTS)
```python
# backend/services/llm_providers/main_image_generation.py

def generate_image(
    prompt: str,
    options: Optional[Dict[str, Any]] = None,
    user_id: Optional[str] = None
) -> ImageGenerationResult:
    """Generate image with pre-flight validation."""
    # 1. Pre-flight validation
    if user_id:
        validate_image_generation_operations(...)
    
    # 2. Select provider
    provider_name = _select_provider(options.get("provider"))
    provider = _get_provider(provider_name)
    
    # 3. Generate
    result = provider.generate(image_options)
    
    # 4. Track usage
    if user_id and result:
        track_image_usage(...)
    
    return result
```

#### **Proposed Extensions** (REUSABLE PATTERN)
```python
# backend/services/llm_providers/main_image_generation.py

# REUSE: Common validation helper
def _validate_image_operation(
    user_id: Optional[str],
    operation_type: str,
    num_operations: int = 1
) -> None:
    """Reusable pre-flight validation for all image operations."""
    if not user_id:
        logger.warning("No user_id provided - skipping validation")
        return
    
    from services.database import get_db
    from services.subscription import PricingService
    from services.subscription.preflight_validator import validate_image_generation_operations
    
    db = next(get_db())
    try:
        pricing_service = PricingService(db)
        validate_image_generation_operations(
            pricing_service=pricing_service,
            user_id=user_id,
            num_images=num_operations
        )
    finally:
        db.close()

# REUSE: Common usage tracking helper
def _track_image_usage(
    user_id: str,
    provider: str,
    model: str,
    operation_type: str,
    result_bytes: bytes,
    cost: float,
    metadata: Optional[Dict[str, Any]] = None
) -> None:
    """Reusable usage tracking for all image operations."""
    # ... (extract from existing generate_image function)

# NEW: Extend for editing operations
def generate_image_edit(
    image_base64: str,
    prompt: str,
    operation: str = "general_edit",
    model: Optional[str] = None,
    options: Optional[Dict[str, Any]] = None,
    user_id: Optional[str] = None
) -> ImageGenerationResult:
    """Generate edited image - REUSES validation and tracking."""
    # 1. Reuse validation
    _validate_image_operation(user_id, "image-edit")
    
    # 2. Get provider (extend to support editing providers)
    provider = _get_edit_provider(model or "wavespeed")
    
    # 3. Generate edit
    result = provider.edit(image_base64, prompt, operation, options)
    
    # 4. Reuse tracking
    if user_id and result:
        _track_image_usage(
            user_id=user_id,
            provider=result.provider,
            model=result.model,
            operation_type="image-edit",
            result_bytes=result.image_bytes,
            cost=result.metadata.get("estimated_cost", 0.0),
            metadata=result.metadata
        )
    
    return result
```

#### **Benefits**
- ✅ **Reuses existing infrastructure** - no duplication
- ✅ **Consistent patterns** - all operations follow same flow
- ✅ **Easy to extend** - add new operations by following pattern
- ✅ **Single source of truth** - validation/tracking in one place

---

### **Pattern 2: Reusable Validation & Tracking Helpers** ✅

#### **Current Implementation** (EXISTS in `main_image_generation.py`)
```python
# Pre-flight validation (lines 58-83)
if user_id:
    db = next(get_db())
    try:
        pricing_service = PricingService(db)
        validate_image_generation_operations(...)
    finally:
        db.close()

# Usage tracking (lines 117-265)
if user_id and result and result.image_bytes:
    # ... tracking logic
```

#### **Proposed Refactoring** (EXTRACT FOR REUSABILITY)
```python
# backend/services/llm_providers/main_image_generation.py

# EXTRACT: Reusable validation function
def _validate_and_track_image_operation(
    user_id: Optional[str],
    operation_type: str,
    provider: str,
    model: str,
    result: Optional[ImageGenerationResult],
    num_operations: int = 1
) -> None:
    """
    REUSABLE helper for validation and tracking.
    Used by all image operation functions.
    """
    # Pre-flight validation
    if user_id:
        _validate_image_operation(user_id, operation_type, num_operations)
    
    # Post-generation tracking
    if user_id and result and result.image_bytes:
        _track_image_usage(
            user_id=user_id,
            provider=provider,
            model=model,
            operation_type=operation_type,
            result_bytes=result.image_bytes,
            cost=result.metadata.get("estimated_cost", 0.0) if result.metadata else 0.0,
            metadata=result.metadata
        )

# REFACTOR: Existing generate_image to use helper
def generate_image(...) -> ImageGenerationResult:
    """Generate image - now uses reusable helpers."""
    # ... provider selection and generation ...
    
    # REUSE: Validation and tracking
    _validate_and_track_image_operation(
        user_id=user_id,
        operation_type="text-to-image",
        provider=provider_name,
        model=result.model,
        result=result
    )
    
    return result
```

#### **Benefits**
- ✅ **DRY Principle** - validation/tracking logic in one place
- ✅ **Consistent behavior** - all operations use same validation
- ✅ **Easy maintenance** - change validation logic once, affects all
- ✅ **Testable** - helpers can be tested independently

---

### **Pattern 3: Extend Provider Pattern for Reusability** ✅

#### **Current Structure** (EXISTS)
```python
# backend/services/llm_providers/image_generation/base.py

class ImageGenerationProvider(Protocol):
    """Protocol for image generation providers."""
    def generate(self, options: ImageGenerationOptions) -> ImageGenerationResult:
        ...

# backend/services/llm_providers/image_generation/wavespeed_provider.py

class WaveSpeedImageProvider(ImageGenerationProvider):
    """WaveSpeed AI image generation provider."""
    SUPPORTED_MODELS = {
        "ideogram-v3-turbo": {...},
        "qwen-image": {...}
    }
    
    def generate(self, options: ImageGenerationOptions) -> ImageGenerationResult:
        # ... implementation
```

#### **Proposed Extension** (REUSE PATTERN)
```python
# backend/services/llm_providers/image_generation/base.py

# EXTEND: Add editing protocol
class ImageEditProvider(Protocol):
    """Protocol for image editing providers."""
    def edit(
        self,
        image_base64: str,
        prompt: str,
        operation: str,
        options: ImageEditOptions
    ) -> ImageGenerationResult:
        ...

# NEW: Reuse WaveSpeed client pattern
# backend/services/llm_providers/image_generation/wavespeed_edit_provider.py

class WaveSpeedEditProvider(ImageEditProvider):
    """WaveSpeed AI image editing provider - REUSES client."""
    
    # REUSE: Same client initialization
    def __init__(self, api_key: Optional[str] = None):
        self.client = WaveSpeedClient(api_key=api_key)  # REUSE
    
    # REUSE: Model registry pattern
    SUPPORTED_MODELS = {
        "qwen-edit": {
            "model_path": "wavespeed-ai/qwen-image/edit",
            "cost": 0.02,
        },
        "step1x-edit": {
            "model_path": "wavespeed-ai/step1x-edit",
            "cost": 0.03,
        },
        # ... 12 editing models
    }
    
    def edit(
        self,
        image_base64: str,
        prompt: str,
        operation: str,
        options: ImageEditOptions
    ) -> ImageGenerationResult:
        """Edit image - REUSES client pattern."""
        model_info = self.SUPPORTED_MODELS.get(options.model)
        if not model_info:
            raise ValueError(f"Unsupported model: {options.model}")
        
        # REUSE: Same client call pattern
        image_bytes = self.client.edit_image(
            model=model_info["model_path"],
            image_base64=image_base64,
            prompt=prompt,
            **options.to_dict()
        )
        
        # REUSE: Same result format
        return ImageGenerationResult(
            image_bytes=image_bytes,
            width=options.width,
            height=options.height,
            provider="wavespeed",
            model=options.model,
            metadata={"cost": model_info["cost"]}
        )
```

#### **Benefits**
- ✅ **Reuses existing protocol pattern** - consistent interface
- ✅ **Reuses WaveSpeedClient** - no duplicate client code
- ✅ **Reuses model registry pattern** - easy to add models
- ✅ **Reuses result format** - consistent return types

---

### **Pattern 4: Reusable Model Registry** (ENHANCE EXISTING)

#### **Current Pattern** (EXISTS in providers)
```python
# WaveSpeedImageProvider.SUPPORTED_MODELS
SUPPORTED_MODELS = {
    "ideogram-v3-turbo": {
        "name": "Ideogram V3 Turbo",
        "cost_per_image": 0.10,
        "max_resolution": (1024, 1024),
    },
    "qwen-image": {...}
}
```

#### **Proposed Enhancement** (CENTRALIZE FOR REUSABILITY)
```python
# backend/services/image_studio/model_registry.py

@dataclass
class ImageModel:
    """Model metadata - REUSES existing provider pattern."""
    id: str
    name: str
    provider: str
    model_path: str
    cost: float
    category: str  # "generation", "editing", "upscaling", "3d", "face-swap"
    capabilities: List[str]
    max_resolution: Optional[tuple[int, int]] = None

class ImageModelRegistry:
    """Centralized registry - AGGREGATES from providers."""
    
    # REUSE: Extract from existing providers
    MODELS: Dict[str, ImageModel] = {
        # Generation (from WaveSpeedImageProvider)
        "ideogram-v3-turbo": ImageModel(
            id="ideogram-v3-turbo",
            name="Ideogram V3 Turbo",
            provider="wavespeed",
            model_path="ideogram-ai/ideogram-v3-turbo",
            cost=0.10,  # From SUPPORTED_MODELS
            category="generation",
            capabilities=["text-to-image"],
        ),
        # Editing (NEW - follows same pattern)
        "qwen-edit": ImageModel(
            id="qwen-edit",
            name="Qwen Image Edit",
            provider="wavespeed",
            model_path="wavespeed-ai/qwen-image/edit",
            cost=0.02,
            category="editing",
            capabilities=["image-edit", "style-transfer"],
        ),
        # ... 40+ models
    }
    
    @classmethod
    def get_model(cls, model_id: str) -> Optional[ImageModel]:
        """Get model by ID - REUSABLE across all services."""
        return cls.MODELS.get(model_id)
    
    @classmethod
    def list_by_category(cls, category: str) -> List[ImageModel]:
        """List models by category - REUSABLE query."""
        return [m for m in cls.MODELS.values() if m.category == category]
    
    @classmethod
    def get_cost(cls, model_id: str) -> float:
        """Get cost for model - REUSABLE cost lookup."""
        model = cls.get_model(model_id)
        return model.cost if model else 0.0
```

#### **Benefits**
- ✅ **Reuses provider model definitions** - single source of truth
- ✅ **Reusable queries** - all services can use same registry
- ✅ **Cost calculation** - centralized cost lookup
- ✅ **Frontend integration** - single endpoint for model list

---

### **Pattern 5: Usage Tracking**

#### **Structure**
```python
# backend/services/llm_providers/main_image_operations.py

def track_image_usage(
    *,
    user_id: str,
    provider: str,
    model_name: str,
    operation_type: str,
    image_bytes: bytes,
    cost_override: Optional[float] = None,
) -> Dict[str, Any]:
    """
    Track subscription usage for image operations.
    Mirrors track_video_usage() pattern.
    """
    from services.database import get_db
    from models.subscription_models import APIProvider, APIUsageLog, UsageSummary
    
    db = next(get_db())
    try:
        pricing_service = PricingService(db)
        current_period = pricing_service.get_current_billing_period(user_id)
        
        # Get or create usage summary
        usage_summary = get_or_create_usage_summary(user_id, current_period)
        
        # Calculate cost
        cost = cost_override or calculate_cost(provider, model_name, operation_type)
        
        # Update usage summary
        update_usage_summary(usage_summary, operation_type, cost)
        
        # Log API usage
        log_api_usage(user_id, provider, model_name, operation_type, cost, image_bytes)
        
        db.commit()
        
        return {
            "previous_calls": previous_count,
            "current_calls": usage_summary.image_calls,
            "cost": cost,
            "total_cost": usage_summary.image_cost,
        }
    finally:
        db.close()
```

#### **Benefits**
- Consistent with video tracking
- Centralized cost calculation
- Automatic usage logging
- Real-time limit checking

---

### **Pattern 6: Service Layer - Reuse Existing Entry Point** ✅

#### **Current Implementation** (MIXED USAGE)
```python
# CreateStudioService - Uses providers directly (NOT using main_image_generation.py)
# Other services (YouTube, Podcast) - Use main_image_generation.py ✅
```

#### **Proposed Refactoring** (REUSE UNIFIED ENTRY)
```python
# backend/services/image_studio/create_service.py

class CreateStudioService:
    """Service for Create Studio - REUSES unified entry point."""
    
    async def generate(
        self,
        request: CreateStudioRequest,
        user_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Generate image - REUSES main_image_generation.py."""
        # REUSE: Existing unified entry point
        from services.llm_providers.main_image_generation import generate_image
        
        # Map request to unified format
        options = {
            "provider": request.provider or "auto",
            "model": request.model,
            "width": request.width,
            "height": request.height,
            "negative_prompt": request.negative_prompt,
            "guidance_scale": request.guidance_scale,
            "steps": request.steps,
            "seed": request.seed,
        }
        
        # REUSE: Call unified entry point
        results = []
        for i in range(request.num_variations):
            result = generate_image(
                prompt=request.prompt,
                options=options,
                user_id=user_id
            )
            results.append({
                "image_bytes": result.image_bytes,
                "width": result.width,
                "height": result.height,
                "model": result.model,
                "metadata": result.metadata,
            })
        
        return {
            "success": True,
            "results": results,
            "cost": sum(r["metadata"].get("estimated_cost", 0) for r in results),
        }
```

#### **Benefits**
- ✅ **Reuses existing unified entry** - no duplicate validation/tracking
- ✅ **Consistent behavior** - all services use same entry point
- ✅ **Thin service layer** - services focus on business logic
- ✅ **Easy to maintain** - changes in entry point affect all services

---

## 🏗️ Implementation Structure (REUSE EXISTING)

### **File Organization** (EXTEND, DON'T DUPLICATE)

```
backend/services/
├── llm_providers/
│   ├── main_image_generation.py      ← EXISTS - EXTEND for new operations
│   │   ✅ generate_image()            (text-to-image)
│   │   ✅ generate_character_image()  (character consistency)
│   │   🆕 generate_image_edit()        (editing operations)
│   │   🆕 generate_image_upscale()     (upscaling)
│   │   🆕 generate_image_to_3d()       (3D generation)
│   │   🆕 generate_face_swap()          (face swapping)
│   │   🆕 generate_image_translate()    (translation)
│   │
│   │   # REUSABLE HELPERS (extract from existing)
│   │   🆕 _validate_image_operation()   (extract validation)
│   │   🆕 _track_image_operation_usage() (extract tracking)
│   │
│   ├── main_video_generation.py      ← Reference pattern
│   │
│   └── image_generation/              ← EXISTS - EXTEND
│       ├── __init__.py                ✅ Exports providers
│       ├── base.py                    ✅ Protocol (EXISTS)
│       │   - ImageGenerationOptions
│       │   - ImageGenerationResult
│       │   - ImageGenerationProvider (Protocol)
│       │   🆕 ImageEditProvider (Protocol)
│       │   🆕 ImageUpscaleProvider (Protocol)
│       │   🆕 Image3DProvider (Protocol)
│       │
│       ├── wavespeed_provider.py      ✅ EXISTS - EXTEND
│       │   - WaveSpeedImageProvider
│       │   🆕 WaveSpeedEditProvider
│       │   🆕 WaveSpeedUpscaleProvider
│       │   🆕 WaveSpeed3DProvider
│       │   🆕 WaveSpeedFaceSwapProvider
│       │
│       ├── stability_provider.py      ✅ EXISTS
│       ├── hf_provider.py             ✅ EXISTS
│       └── gemini_provider.py         ✅ EXISTS
│
├── image_studio/
│   ├── studio_manager.py             ✅ EXISTS (orchestrator)
│   ├── create_service.py              ⚠️ REFACTOR: Use main_image_generation
│   ├── edit_service.py                ⚠️ REFACTOR: Use main_image_generation
│   ├── upscale_service.py             ⚠️ REFACTOR: Use main_image_generation
│   ├── transform_service.py           ✅ Uses main_video_generation
│   ├── three_d_service.py             🆕 NEW: Uses main_image_generation
│   ├── face_swap_service.py           🆕 NEW: Uses main_image_generation
│   └── model_registry.py               🆕 NEW: Centralized registry
│
└── subscription/
    └── preflight_validator.py         ✅ EXISTS - REUSE
        - validate_image_generation_operations()
```

### **Key Reusability Principles**

1. **Extend, Don't Duplicate**
   - ✅ Extend `main_image_generation.py` (don't create new file)
   - ✅ Extend `ImageGenerationProvider` protocol (don't create new base)
   - ✅ Reuse `WaveSpeedClient` (don't duplicate client code)

2. **Extract Common Logic**
   - ✅ Extract validation into reusable helper
   - ✅ Extract tracking into reusable helper
   - ✅ Extract cost calculation into reusable helper

3. **Consistent Patterns**
   - ✅ All operations follow same function signature pattern
   - ✅ All operations use same validation/tracking helpers
   - ✅ All providers follow same protocol pattern

---

## 🔄 Implementation Strategy (REUSE EXISTING)

### **Phase 1: Extract Reusable Helpers** (Week 1)
1. ✅ **Extract validation helper** from `generate_image()` → `_validate_image_operation()`
2. ✅ **Extract tracking helper** from `generate_image()` → `_track_image_operation_usage()`
3. ✅ **Refactor existing functions** to use extracted helpers
4. ✅ **Test** - ensure existing functionality unchanged

### **Phase 2: Extend for Editing** (Week 2)
1. ✅ **Add `ImageEditProvider` protocol** to `base.py`
2. ✅ **Create `WaveSpeedEditProvider`** following existing provider pattern
3. ✅ **Add `generate_image_edit()`** to `main_image_generation.py` (reuses helpers)
4. ✅ **Refactor `EditStudioService`** to use unified entry point

### **Phase 3: Extend for Upscaling** (Week 3)
1. ✅ **Add `ImageUpscaleProvider` protocol** to `base.py`
2. ✅ **Create `WaveSpeedUpscaleProvider`** (reuses WaveSpeedClient)
3. ✅ **Add `generate_image_upscale()`** (reuses validation/tracking)
4. ✅ **Refactor `UpscaleStudioService`** to use unified entry

### **Phase 4: Extend for 3D & Specialized** (Week 4-5)
1. ✅ **Add `Image3DProvider` protocol**
2. ✅ **Create `WaveSpeed3DProvider`** (reuses client pattern)
3. ✅ **Add `generate_image_to_3d()`** (reuses helpers)
4. ✅ **Add face swap, translation** following same pattern
5. ✅ **Create new services** (3D, Face Swap) using unified entry

### **Phase 5: Model Registry** (Week 6)
1. ✅ **Create `model_registry.py`** aggregating from providers
2. ✅ **Update providers** to register models in central registry
3. ✅ **Add API endpoint** for model list (frontend integration)
4. ✅ **Update cost estimation** to use registry

### **Key Principles**
- ✅ **Reuse existing code** - don't duplicate
- ✅ **Extract common logic** - DRY principle
- ✅ **Follow existing patterns** - consistency
- ✅ **Test incrementally** - ensure no regressions

---

## 📋 Reusable Code Examples

### **Example 1: Adding a New Editing Model** (REUSES PATTERNS)

```python
# 1. Add to WaveSpeedEditProvider (REUSES existing pattern)
# backend/services/llm_providers/image_generation/wavespeed_edit_provider.py

class WaveSpeedEditProvider(ImageEditProvider):
    SUPPORTED_MODELS = {
        # ... existing models ...
        "new-edit-model": {  # 🆕 NEW MODEL
            "model_path": "wavespeed-ai/new-edit-model",
            "cost": 0.05,
            "max_resolution": (2048, 2048),
        }
    }
    
    def edit(self, image_base64: str, prompt: str, ...):
        # REUSES: Same client call pattern
        model_info = self.SUPPORTED_MODELS.get(options.model)
        image_bytes = self.client.edit_image(
            model=model_info["model_path"],
            image_base64=image_base64,
            prompt=prompt,
            **options.to_dict()
        )
        # REUSES: Same result format
        return ImageGenerationResult(...)

# 2. Register in model registry (REUSES registry pattern)
# backend/services/image_studio/model_registry.py
ImageModelRegistry.MODELS["new-edit-model"] = ImageModel(
    id="new-edit-model",
    name="New Edit Model",
    provider="wavespeed",
    model_path="wavespeed-ai/new-edit-model",
    cost=0.05,  # From provider SUPPORTED_MODELS
    category="editing",
    capabilities=["image-edit"],
)

# 3. Use in service (REUSES unified entry)
# backend/services/image_studio/edit_service.py
from services.llm_providers.main_image_generation import generate_image_edit

result = generate_image_edit(
    image_base64=image,
    prompt=prompt,
    model="new-edit-model",  # 🆕 Just specify model ID
    user_id=user_id,
)
# ✅ Validation, tracking, error handling all handled automatically
```

### **Example 2: Adding a New Operation Type** (REUSES HELPERS)

```python
# In main_image_generation.py (EXTEND existing file)

def generate_face_swap(
    source_image_base64: str,
    target_image_base64: str,
    model: str = "wavespeed-ai/image-face-swap",
    options: Optional[Dict[str, Any]] = None,
    user_id: Optional[str] = None
) -> ImageGenerationResult:
    """
    Face swap operation - REUSES validation and tracking helpers.
    """
    # 1. REUSE: Validation helper
    _validate_image_operation(user_id, "face-swap")
    
    # 2. Get provider (REUSES provider pattern)
    provider = _get_face_swap_provider(model)
    
    # 3. Perform operation
    result = provider.face_swap(
        source_image_base64=source_image_base64,
        target_image_base64=target_image_base64,
        model=model,
        options=options or {}
    )
    
    # 4. REUSE: Tracking helper
    if user_id and result:
        _track_image_operation_usage(
            user_id=user_id,
            provider=result.provider,
            model=result.model,
            operation_type="face-swap",
            result_bytes=result.image_bytes,
            cost=result.metadata.get("estimated_cost", 0.0),
            metadata=result.metadata
        )
    
    return result
```

### **Example 3: Refactoring Existing Service** (REUSE UNIFIED ENTRY)

```python
# BEFORE: CreateStudioService uses providers directly
class CreateStudioService:
    async def generate(self, request, user_id):
        # ... validation logic ...
        provider = self._get_provider_instance(provider_name)
        result = provider.generate(options)
        # ... tracking logic ...
        return result

# AFTER: CreateStudioService REUSES unified entry
class CreateStudioService:
    async def generate(self, request, user_id):
        # REUSE: Unified entry point (validation + tracking included)
        from services.llm_providers.main_image_generation import generate_image
        
        results = []
        for i in range(request.num_variations):
            result = generate_image(  # ✅ All validation/tracking handled
                prompt=request.prompt,
                options={...},
                user_id=user_id
            )
            results.append(result)
        
        return {"results": results}
```

---

## ✅ Benefits of Reusable Architecture

1. **✅ Reuses Existing Code**: Builds on `main_image_generation.py` (no duplication)
2. **✅ DRY Principle**: Validation and tracking extracted into reusable helpers
3. **✅ Consistent Patterns**: All operations follow same proven pattern
4. **✅ Easy to Extend**: Add new operations by following existing pattern
5. **✅ Single Source of Truth**: Model registry aggregates from providers
6. **✅ Maintainable**: Changes in helpers affect all operations
7. **✅ Testable**: Helpers can be tested independently
8. **✅ Backward Compatible**: Existing code continues to work

---

## 🎯 Next Steps

1. **✅ Review existing `main_image_generation.py`** - understand current implementation
2. **✅ Extract reusable helpers** - validation and tracking functions
3. **✅ Extend for editing operations** - add `generate_image_edit()` following pattern
4. **✅ Create model registry** - aggregate models from all providers
5. **✅ Refactor services** - make them use unified entry point
6. **✅ Add new operations** - 3D, face swap, translation following same pattern

## 📝 Implementation Checklist

### **Reusability Focus**
- [ ] Extract `_validate_image_operation()` helper from existing code
- [ ] Extract `_track_image_operation_usage()` helper from existing code
- [ ] Refactor `generate_image()` to use extracted helpers
- [ ] Refactor `generate_character_image()` to use extracted helpers
- [ ] Add `generate_image_edit()` using same helpers
- [ ] Add `generate_image_upscale()` using same helpers
- [ ] Add `generate_image_to_3d()` using same helpers
- [ ] Create `ImageModelRegistry` aggregating from providers
- [ ] Refactor `CreateStudioService` to use unified entry
- [ ] Refactor `EditStudioService` to use unified entry
- [ ] All new operations follow same pattern

---

## 🎯 Reusability Implementation Roadmap

### **Phase 1: Extract Reusable Helpers** (Week 1)
**Goal**: Extract common logic from existing code

1. ✅ **Extract `_validate_image_operation()`** from `generate_image()` (lines 58-83)
2. ✅ **Extract `_track_image_operation_usage()`** from `generate_image()` (lines 117-265)
3. ✅ **Refactor existing functions** to use extracted helpers
4. ✅ **Test** - ensure no regressions

### **Phase 2: Extend for Editing** (Week 2)
**Goal**: Add editing operations reusing patterns

1. ✅ **Add `ImageEditProvider` protocol** to `base.py` (reuses protocol pattern)
2. ✅ **Create `WaveSpeedEditProvider`** (reuses WaveSpeedClient, model registry pattern)
3. ✅ **Add `generate_image_edit()`** to `main_image_generation.py` (reuses helpers)
4. ✅ **Refactor `EditStudioService`** to use unified entry

### **Phase 3: Extend for Other Operations** (Week 3-4)
**Goal**: Add upscaling, 3D, face swap following same pattern

- Same approach as Phase 2 for each operation type

### **Phase 4: Model Registry** (Week 5)
**Goal**: Centralize model information

- Aggregate models from all providers
- Single source of truth for cost, capabilities, etc.

---

## 📚 Related Documentation

- [Image Studio Enhancement Proposal](docs/IMAGE_STUDIO_ENHANCEMENT_PROPOSAL.md) - **Updated with reusability focus**
- [Code Patterns Reference](docs/IMAGE_STUDIO_CODE_PATTERNS_REFERENCE.md) - **Reusability patterns**
- [WaveSpeed Models Reference](docs/IMAGE_STUDIO_WAVESPEED_MODELS_REFERENCE.md)
- [Image Studio Implementation Review](docs/IMAGE_STUDIO_IMPLEMENTATION_REVIEW.md)
- [Video Studio Implementation](backend/services/llm_providers/main_video_generation.py) - Reference pattern

---

*Document Version: 2.0*  
*Last Updated: Current Session*  
*Status: Architecture Proposal - Reusability Focus*  
*Key Principle: Extend existing `main_image_generation.py`, don't duplicate*