main.py

"""
Byte Dream - Main Application Interface
Simple Python API for image generation
"""

from bytedream.generator import ByteDreamGenerator
from bytedream.utils import (
    load_image,
    save_image,
    resize_image,
    create_image_grid,
)
from typing import Optional, List
from PIL import Image


class ByteDreamApp:
    """
    High-level application interface for Byte Dream
    Simplifies common tasks like image generation and batch processing
    """
    
    def __init__(
        self,
        model_path: Optional[str] = None,
        device: str = "cpu",
        verbose: bool = True,
    ):
        """
        Initialize Byte Dream application
        
        Args:
            model_path: Path to model weights
            device: Device to run on
            verbose: Enable verbose output
        """
        self.verbose = verbose
        
        if self.verbose:
            print("Initializing Byte Dream Application...")
        
        self.generator = ByteDreamGenerator(
            model_path=model_path,
            config_path="config.yaml",
            device=device,
        )
        
        if self.verbose:
            print("✓ Application ready!")
    
    def generate(
        self,
        prompt: str,
        output_path: str = "output.png",
        negative_prompt: Optional[str] = None,
        width: int = 512,
        height: int = 512,
        steps: int = 50,
        guidance: float = 7.5,
        seed: Optional[int] = None,
        save: bool = True,
    ) -> Image.Image:
        """
        Generate image from prompt and optionally save to file
        
        Args:
            prompt: Text description
            output_path: Where to save the image
            negative_prompt: What to avoid
            width: Image width
            height: Image height
            steps: Inference steps
            guidance: Guidance scale
            seed: Random seed
            save: Whether to save to file
            
        Returns:
            Generated PIL Image
        """
        # Generate image
        image = self.generator.generate(
            prompt=prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            num_inference_steps=steps,
            guidance_scale=guidance,
            seed=seed,
        )
        
        # Save if requested
        if save:
            save_image(image, output_path)
            if self.verbose:
                print(f"✓ Image saved to: {output_path}")
        
        return image
    
    def generate_multiple(
        self,
        prompts: List[str],
        output_dir: str = "./outputs",
        negative_prompt: Optional[str] = None,
        width: int = 512,
        height: int = 512,
        steps: int = 50,
        guidance: float = 7.5,
        seeds: Optional[List[int]] = None,
        create_grid: bool = True,
    ) -> List[Image.Image]:
        """
        Generate multiple images from prompts
        
        Args:
            prompts: List of prompts
            output_dir: Directory to save images
            negative_prompt: Negative prompt for all
            width: Image width
            height: Image height
            steps: Inference steps
            guidance: Guidance scale
            seeds: Seeds for each image
            create_grid: Create grid of all images
            
        Returns:
            List of generated images
        """
        from pathlib import Path
        
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        
        images = []
        
        for i, prompt in enumerate(prompts):
            print(f"\n{'='*60}")
            print(f"Generating image {i+1}/{len(prompts)}")
            print(f"{'='*60}")
            
            seed = seeds[i] if seeds else None
            
            image = self.generate(
                prompt=prompt,
                output_path=str(output_path / f"image_{i+1:03d}.png"),
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                steps=steps,
                guidance=guidance,
                seed=seed,
                save=True,
            )
            
            images.append(image)
        
        # Create grid
        if create_grid and len(images) > 1:
            grid = create_image_grid(images)
            grid_path = output_path / "grid.png"
            grid.save(grid_path)
            print(f"\n✓ Grid saved to: {grid_path}")
        
        return images
    
    def img2img(
        self,
        input_image_path: str,
        prompt: str,
        output_path: str = "output_img2img.png",
        strength: float = 0.75,
        negative_prompt: Optional[str] = None,
        steps: int = 50,
        guidance: float = 7.5,
        seed: Optional[int] = None,
    ) -> Image.Image:
        """
        Image-to-image transformation (placeholder for future implementation)
        
        Args:
            input_image_path: Input image path
            prompt: Transformation prompt
            output_path: Output path
            strength: How much to transform (0-1)
            negative_prompt: Negative prompt
            steps: Inference steps
            guidance: Guidance scale
            seed: Random seed
            
        Returns:
            Transformed image
        """
        print("⚠ img2img functionality will be available in a future update")
        print("  For now, using text-to-image generation only")
        
        # For now, just generate from prompt
        return self.generate(
            prompt=prompt,
            output_path=output_path,
            negative_prompt=negative_prompt,
            steps=steps,
            guidance=guidance,
            seed=seed,
        )
    
    def info(self):
        """Print model information"""
        info = self.generator.get_model_info()
        
        print("\n" + "="*60)
        print("Byte Dream Model Information")
        print("="*60)
        for key, value in info.items():
            print(f"{key.replace('_', ' ').title()}: {value}")
        print("="*60)


def demo():
    """Run a quick demo"""
    print("\n" + "="*60)
    print("Byte Dream - Quick Demo")
    print("="*60)
    
    app = ByteDreamApp(device="cpu", verbose=True)
    
    # Demo prompts
    prompts = [
        "A beautiful sunset over mountains, digital art, vibrant colors",
        "Cyberpunk city at night with neon lights, futuristic",
        "Fantasy landscape with castle and waterfall, epic",
    ]
    
    print("\nGenerating sample images...")
    
    images = app.generate_multiple(
        prompts=prompts,
        output_dir="./demo_outputs",
        steps=30,  # Fewer steps for demo
        guidance=7.5,
        create_grid=True,
    )
    
    print(f"\n✓ Demo complete! Generated {len(images)} images")
    print("  Check ./demo_outputs/ for results")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Byte Dream Application")
    parser.add_argument("--demo", action="store_true", help="Run demo")
    args = parser.parse_args()
    
    if args.demo:
        demo()
    else:
        # Interactive mode
        app = ByteDreamApp()
        
        print("\nByte Dream Interactive Mode")
        print("Type 'quit' to exit\n")
        
        while True:
            prompt = input("Prompt: ").strip()
            
            if prompt.lower() in ['quit', 'exit', 'q']:
                break
            
            if not prompt:
                continue
            
            try:
                image = app.generate(
                    prompt=prompt,
                    output_path=f"output_{len(prompt)}.png",
                )
                print("✓ Image generated!\n")
            except Exception as e:
                print(f"Error: {e}\n")