Create pipeline.py

pipeline.py

@@ -0,0 +1,551 @@
+#!/usr/bin/env python3
+"""
+Trouter-Imagine-1 Complete Pipeline
+Apache 2.0 License
+
+This file provides a complete, ready-to-use pipeline for text-to-image generation.
+It includes all necessary components and can be used immediately for generating images.
+
+This is the MAIN FILE for using the model - simple and powerful.
+"""
+
+import torch
+from diffusers import (
+    StableDiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    DDIMScheduler
+)
+from PIL import Image
+import os
+from typing import List, Optional, Union, Dict
+import warnings
+import logging
+from pathlib import Path
+import json
+from datetime import datetime
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class TrouterImagePipeline:
+    """
+    Complete ready-to-use pipeline for Trouter-Imagine-1
+    
+    This is the main class you should use for image generation.
+    It's simple, powerful, and handles everything automatically.
+    
+    Example:
+        >>> pipeline = TrouterImagePipeline()
+        >>> image = pipeline("a beautiful sunset")
+        >>> image.save("sunset.png")
+    """
+    
+    # Default base model (you can change this to your custom model once trained)
+    DEFAULT_MODEL = "runwayml/stable-diffusion-v1-5"
+    
+    # Can also use these alternatives:
+    # "stabilityai/stable-diffusion-2-1"
+    # "stabilityai/stable-diffusion-xl-base-1.0"
+    
+    def __init__(
+        self,
+        model_id: Optional[str] = None,
+        device: Optional[str] = None,
+        torch_dtype: torch.dtype = torch.float16,
+        use_safetensors: bool = True,
+        enable_optimizations: bool = True
+    ):
+        """
+        Initialize the Trouter-Imagine-1 pipeline
+        
+        Args:
+            model_id: Model to use (defaults to Stable Diffusion 1.5)
+            device: Device to use (auto-detected if None)
+            torch_dtype: Model precision (float16 for speed, float32 for quality)
+            use_safetensors: Use safetensors format (recommended)
+            enable_optimizations: Enable memory optimizations
+        """
+        # Auto-detect device
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+                logger.info(f"CUDA available: {torch.cuda.get_device_name(0)}")
+            elif torch.backends.mps.is_available():
+                device = "mps"
+                logger.info("Using Apple Silicon (MPS)")
+            else:
+                device = "cpu"
+                logger.warning("No GPU detected, using CPU (will be slow)")
+        
+        self.device = device
+        self.dtype = torch_dtype
+        self.model_id = model_id or self.DEFAULT_MODEL
+        
+        logger.info(f"Initializing Trouter-Imagine-1 Pipeline")
+        logger.info(f"Model: {self.model_id}")
+        logger.info(f"Device: {self.device}")
+        logger.info(f"Precision: {self.dtype}")
+        
+        # Load pipeline
+        self._load_pipeline(use_safetensors)
+        
+        # Apply optimizations
+        if enable_optimizations:
+            self._optimize()
+        
+        # Default settings
+        self.default_negative = "blurry, low quality, distorted, deformed, ugly, bad anatomy, watermark, signature, text"
+        
+        logger.info("✓ Pipeline ready!")
+    
+    def _load_pipeline(self, use_safetensors: bool):
+        """Load the diffusion pipeline"""
+        try:
+            self.pipe = StableDiffusionPipeline.from_pretrained(
+                self.model_id,
+                torch_dtype=self.dtype,
+                use_safetensors=use_safetensors,
+                safety_checker=None,  # Disable for flexibility
+                requires_safety_checker=False
+            )
+            
+            # Move to device
+            self.pipe = self.pipe.to(self.device)
+            
+            # Set better scheduler by default
+            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+                self.pipe.scheduler.config
+            )
+            
+            logger.info("✓ Model loaded successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            raise
+    
+    def _optimize(self):
+        """Apply memory and speed optimizations"""
+        logger.info("Applying optimizations...")
+        
+        try:
+            # Memory optimizations
+            self.pipe.enable_attention_slicing()
+            self.pipe.enable_vae_slicing()
+            logger.info("  ✓ Memory optimizations enabled")
+        except Exception as e:
+            logger.warning(f"  ⚠ Memory optimization failed: {e}")
+        
+        # Try xformers for even better performance
+        try:
+            self.pipe.enable_xformers_memory_efficient_attention()
+            logger.info("  ✓ xformers enabled (faster generation)")
+        except Exception:
+            logger.info("  ℹ xformers not available (this is fine)")
+        
+        # Model CPU offload for very limited VRAM
+        # Uncomment if you have < 6GB VRAM:
+        # self.pipe.enable_model_cpu_offload()
+    
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 7.5,
+        num_images: int = 1,
+        seed: Optional[int] = None,
+        return_dict: bool = False
+    ) -> Union[Image.Image, List[Image.Image], Dict]:
+        """
+        Generate images from text prompt
+        
+        Args:
+            prompt: Text description or list of descriptions
+            negative_prompt: What to avoid (uses default if None)
+            width: Image width (must be multiple of 8)
+            height: Image height (must be multiple of 8)
+            num_inference_steps: Quality (20=fast, 30=balanced, 50=quality)
+            guidance_scale: Prompt adherence (7.5 is good default)
+            num_images: Number of images to generate
+            seed: Random seed for reproducibility
+            return_dict: Return dictionary with metadata
+        
+        Returns:
+            Generated image(s) or dictionary with images and metadata
+        """
+        # Use default negative prompt if none provided
+        if negative_prompt is None:
+            negative_prompt = self.default_negative
+        
+        # Set seed if provided
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+        
+        # Validate dimensions
+        if width % 8 != 0:
+            width = (width // 8) * 8
+            logger.warning(f"Width adjusted to {width} (must be multiple of 8)")
+        if height % 8 != 0:
+            height = (height // 8) * 8
+            logger.warning(f"Height adjusted to {height} (must be multiple of 8)")
+        
+        # Generate
+        logger.info(f"Generating: {prompt[:100]}...")
+        
+        try:
+            with torch.autocast(self.device) if self.device == "cuda" else torch.no_grad():
+                output = self.pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=guidance_scale,
+                    num_images_per_prompt=num_images,
+                    generator=generator
+                )
+            
+            images = output.images
+            logger.info(f"✓ Generated {len(images)} image(s)")
+            
+            if return_dict:
+                return {
+                    'images': images,
+                    'prompt': prompt,
+                    'negative_prompt': negative_prompt,
+                    'width': width,
+                    'height': height,
+                    'steps': num_inference_steps,
+                    'guidance': guidance_scale,
+                    'seed': seed
+                }
+            
+            return images[0] if len(images) == 1 else images
+            
+        except torch.cuda.OutOfMemoryError:
+            logger.error("GPU out of memory! Try:")
+            logger.error("  1. Reduce resolution (e.g., 512x512 instead of 1024x1024)")
+            logger.error("  2. Reduce num_images")
+            logger.error("  3. Close other applications")
+            raise
+        except Exception as e:
+            logger.error(f"Generation failed: {e}")
+            raise
+    
+    def generate_batch(
+        self,
+        prompts: List[str],
+        output_dir: str = "./outputs",
+        **kwargs
+    ) -> List[Image.Image]:
+        """
+        Generate multiple images from different prompts
+        
+        Args:
+            prompts: List of text prompts
+            output_dir: Directory to save images
+            **kwargs: Additional generation parameters
+        
+        Returns:
+            List of generated images
+        """
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+        
+        images = []
+        logger.info(f"Generating batch of {len(prompts)} images...")
+        
+        for i, prompt in enumerate(prompts):
+            logger.info(f"  [{i+1}/{len(prompts)}] {prompt[:50]}...")
+            
+            image = self(prompt, **kwargs)
+            images.append(image)
+            
+            # Save
+            filename = output_path / f"image_{i:04d}.png"
+            image.save(filename)
+            logger.info(f"    ✓ Saved to {filename}")
+        
+        logger.info(f"✓ Batch complete! {len(images)} images in {output_dir}")
+        return images
+    
+    def generate_variations(
+        self,
+        prompt: str,
+        num_variations: int = 4,
+        **kwargs
+    ) -> List[Image.Image]:
+        """
+        Generate variations of the same prompt (different seeds)
+        
+        Args:
+            prompt: Text prompt
+            num_variations: Number of variations
+            **kwargs: Additional generation parameters
+        
+        Returns:
+            List of image variations
+        """
+        logger.info(f"Generating {num_variations} variations...")
+        
+        images = []
+        for i in range(num_variations):
+            seed = torch.randint(0, 2**32, (1,)).item()
+            image = self(prompt, seed=seed, **kwargs)
+            images.append(image)
+            logger.info(f"  ✓ Variation {i+1}/{num_variations}")
+        
+        return images
+    
+    def set_scheduler(self, scheduler_name: str):
+        """
+        Change the diffusion scheduler
+        
+        Args:
+            scheduler_name: 'dpm' (fast), 'euler' (creative), 'ddim' (stable)
+        """
+        schedulers = {
+            'dpm': DPMSolverMultistepScheduler,
+            'euler': EulerAncestralDiscreteScheduler,
+            'ddim': DDIMScheduler,
+        }
+        
+        if scheduler_name.lower() not in schedulers:
+            logger.warning(f"Unknown scheduler: {scheduler_name}")
+            return
+        
+        scheduler_class = schedulers[scheduler_name.lower()]
+        self.pipe.scheduler = scheduler_class.from_config(
+            self.pipe.scheduler.config
+        )
+        logger.info(f"✓ Scheduler changed to {scheduler_name}")
+    
+    def save_pipeline(self, save_path: str):
+        """Save the complete pipeline"""
+        self.pipe.save_pretrained(save_path)
+        logger.info(f"✓ Pipeline saved to {save_path}")
+    
+    def get_config(self) -> Dict:
+        """Get current pipeline configuration"""
+        return {
+            'model_id': self.model_id,
+            'device': str(self.device),
+            'dtype': str(self.dtype),
+            'scheduler': self.pipe.scheduler.__class__.__name__,
+            'default_negative_prompt': self.default_negative
+        }
+
+
+# ============================================================================
+# CONVENIENCE FUNCTIONS
+# ============================================================================
+
+def quick_generate(
+    prompt: str,
+    output_path: str = "output.png",
+    quality: str = "balanced",
+    **kwargs
+) -> Image.Image:
+    """
+    Quick one-line image generation
+    
+    Args:
+        prompt: What to generate
+        output_path: Where to save
+        quality: 'draft' (fast), 'balanced', 'high', 'ultra'
+        **kwargs: Additional parameters
+    
+    Returns:
+        Generated image
+    
+    Example:
+        >>> quick_generate("a cat in a hat", "cat.png")
+    """
+    quality_presets = {
+        'draft': {'num_inference_steps': 15, 'width': 512, 'height': 512},
+        'balanced': {'num_inference_steps': 30, 'width': 512, 'height': 512},
+        'high': {'num_inference_steps': 40, 'width': 768, 'height': 768},
+        'ultra': {'num_inference_steps': 50, 'width': 1024, 'height': 1024}
+    }
+    
+    settings = quality_presets.get(quality, quality_presets['balanced'])
+    settings.update(kwargs)
+    
+    pipeline = TrouterImagePipeline()
+    image = pipeline(prompt, **settings)
+    image.save(output_path)
+    
+    logger.info(f"✓ Image saved to {output_path}")
+    return image
+
+
+def batch_from_file(
+    prompts_file: str,
+    output_dir: str = "./outputs",
+    **kwargs
+) -> List[Image.Image]:
+    """
+    Generate images from prompts in a text file
+    
+    Args:
+        prompts_file: Text file with one prompt per line
+        output_dir: Where to save images
+        **kwargs: Generation parameters
+    
+    Returns:
+        List of generated images
+    """
+    with open(prompts_file, 'r') as f:
+        prompts = [line.strip() for line in f if line.strip()]
+    
+    pipeline = TrouterImagePipeline()
+    return pipeline.generate_batch(prompts, output_dir, **kwargs)
+
+
+# ============================================================================
+# PRESETS AND STYLES
+# ============================================================================
+
+STYLE_PRESETS = {
+    'photorealistic': {
+        'prompt_suffix': ', professional photography, photorealistic, 4k, highly detailed',
+        'negative_prompt': 'cartoon, anime, painting, illustration, low quality, blurry',
+        'guidance_scale': 8.5
+    },
+    'artistic': {
+        'prompt_suffix': ', digital art, concept art, detailed illustration',
+        'negative_prompt': 'photograph, realistic, blurry, low quality',
+        'guidance_scale': 7.0
+    },
+    'anime': {
+        'prompt_suffix': ', anime style, manga, cel shaded, vibrant colors',
+        'negative_prompt': 'realistic, 3d, photograph, blurry, low quality',
+        'guidance_scale': 7.5
+    },
+    'oil_painting': {
+        'prompt_suffix': ', oil painting, painterly, artistic, brushstrokes',
+        'negative_prompt': 'photograph, digital, 3d render, blurry',
+        'guidance_scale': 7.5
+    },
+    'cinematic': {
+        'prompt_suffix': ', cinematic lighting, film still, dramatic, movie scene',
+        'negative_prompt': 'amateur, low quality, poor lighting, blurry',
+        'guidance_scale': 8.0
+    }
+}
+
+
+def generate_with_style(
+    prompt: str,
+    style: str = 'photorealistic',
+    output_path: str = "styled_output.png",
+    **kwargs
+) -> Image.Image:
+    """
+    Generate image with predefined style preset
+    
+    Args:
+        prompt: Base prompt
+        style: Style preset name
+        output_path: Where to save
+        **kwargs: Additional parameters
+    
+    Returns:
+        Generated image
+    """
+    if style not in STYLE_PRESETS:
+        logger.warning(f"Unknown style: {style}, using photorealistic")
+        style = 'photorealistic'
+    
+    preset = STYLE_PRESETS[style]
+    
+    # Apply style
+    full_prompt = prompt + preset['prompt_suffix']
+    kwargs['negative_prompt'] = preset['negative_prompt']
+    kwargs['guidance_scale'] = preset['guidance_scale']
+    
+    pipeline = TrouterImagePipeline()
+    image = pipeline(full_prompt, **kwargs)
+    image.save(output_path)
+    
+    logger.info(f"✓ {style.title()} style image saved to {output_path}")
+    return image
+
+
+# ============================================================================
+# MAIN - COMMAND LINE INTERFACE
+# ============================================================================
+
+def main():
+    """Simple command line interface"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Trouter-Imagine-1 Image Generator")
+    parser.add_argument("prompt", type=str, help="Text prompt for generation")
+    parser.add_argument("--output", "-o", type=str, default="output.png",
+                       help="Output file path")
+    parser.add_argument("--quality", "-q", type=str, default="balanced",
+                       choices=['draft', 'balanced', 'high', 'ultra'],
+                       help="Quality preset")
+    parser.add_argument("--style", "-s", type=str,
+                       choices=list(STYLE_PRESETS.keys()),
+                       help="Style preset")
+    parser.add_argument("--seed", type=int, help="Random seed")
+    parser.add_argument("--width", type=int, default=512, help="Image width")
+    parser.add_argument("--height", type=int, default=512, help="Image height")
+    parser.add_argument("--steps", type=int, default=30, help="Inference steps")
+    parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale")
+    parser.add_argument("--negative", type=str, help="Negative prompt")
+    
+    args = parser.parse_args()
+    
+    kwargs = {
+        'width': args.width,
+        'height': args.height,
+        'num_inference_steps': args.steps,
+        'guidance_scale': args.guidance,
+        'seed': args.seed
+    }
+    
+    if args.negative:
+        kwargs['negative_prompt'] = args.negative
+    
+    if args.style:
+        generate_with_style(args.prompt, args.style, args.output, **kwargs)
+    else:
+        quick_generate(args.prompt, args.output, args.quality, **kwargs)
+
+
+if __name__ == "__main__":
+    print("="*70)
+    print("TROUTER-IMAGINE-1 IMAGE GENERATION PIPELINE")
+    print("Apache 2.0 License")
+    print("="*70)
+    print()
+    print("Quick Start Examples:")
+    print()
+    print("  # Python:")
+    print("  from pipeline import TrouterImagePipeline")
+    print("  pipeline = TrouterImagePipeline()")
+    print("  image = pipeline('a beautiful sunset over mountains')")
+    print("  image.save('sunset.png')")
+    print()
+    print("  # Command line:")
+    print("  python pipeline.py 'a cat in a hat' --output cat.png")
+    print("  python pipeline.py 'portrait' --style photorealistic --quality high")
+    print()
+    print("="*70)
+    print()
+    
+    # Run CLI if arguments provided
+    import sys
+    if len(sys.argv) > 1:
+        main()