ALwrity/docs/image studio/IMAGE_STUDIO_ARCHITECTURE_PROPOSAL.md

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)

# 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)

# 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)

# 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)

# 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)

# 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)

# 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)

# 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)

# 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

# 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)

# CreateStudioService - Uses providers directly (NOT using main_image_generation.py)
# Other services (YouTube, Podcast) - Use main_image_generation.py ✅

Proposed Refactoring (REUSE UNIFIED ENTRY)

# 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)

# 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)

# 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)

# 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 - Updated with reusability focus
• Code Patterns Reference - Reusability patterns
• WaveSpeed Models Reference
• Image Studio Implementation Review
• Video Studio Implementation - 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