mock_components.py

import fal_client
from weave_prompt import TextToImageModel, ImageEvaluator, PromptRefiner
from PIL import Image
import numpy as np
from typing import Dict, Any
import os
from fal_image_generator import FalImageGenerator

class MockTextToImageModel(TextToImageModel):
    """Mock text-to-image model for demonstration."""
    
    def __init__(self):
        self.image_generator = FalImageGenerator()
    
    def generate(self, prompt: str, **kwargs) -> Image.Image:
        """Generate an image using the fal image generator."""
        return self.image_generator.generate_image(prompt, **kwargs)

class MockImageEvaluator(ImageEvaluator):
    """Mock image evaluator for demonstration."""
    
    def generate_initial_prompt(self, target_img: Image.Image) -> str:
        """Generate a mock initial prompt."""
        return "A beautiful image with vibrant colors"

    def analyze_differences(self, generated_img: Image.Image, target_img: Image.Image) -> Dict[str, Any]:
        """Mock difference analysis."""
        return {
            "missing_elements": ["texture", "details"],
            "style_differences": ["color intensity", "composition"]
        }

class MockSimilarityMetric:
    """Mock similarity metric that gradually increases."""
    
    def compute(self, generated_img: Image.Image, target_img: Image.Image) -> float:
        """Mock similarity computation that gradually increases."""
        # Randomly increase similarity over time
        return np.random.uniform(0.5, 0.95)

class MockPromptRefiner(PromptRefiner):
    """Mock prompt refiner for demonstration."""
    
    def refine_prompt(self, current_prompt: str, analysis: Dict[str, Any], similarity_score: float) -> str:
        """Mock prompt refinement by adding random modifiers."""
        modifiers = [
            "with more detail",
            "in vibrant colors",
            "with better composition",
            "high quality",
            "masterfully crafted"
        ]
        return f"{current_prompt}, {np.random.choice(modifiers)}"