Stable Diffusion - Criando Arte com IA

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image
import os
from datetime import datetime
import json

class ArtisticImageGenerator:
    def __init__(self):
        print("🎨 Inicializando Gerador de Arte...")
        
        # Configurar device
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"📱 Usando: {self.device}")
        
        # Carregar modelo
        model_id = "runwayml/stable-diffusion-v1-5"
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
            use_safetensors=True
        )
        
        # Otimizar scheduler
        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            self.pipe.scheduler.config
        )
        
        # Mover para device
        self.pipe = self.pipe.to(self.device)
        
        # Otimizações de memória
        if self.device == "cuda":
            self.pipe.enable_model_cpu_offload()
            self.pipe.enable_xformers_memory_efficient_attention()
        
        print("✅ Gerador pronto!")
        
        # Estilos predefinidos
        self.art_styles = {
            "realism": "photorealistic, highly detailed, professional photography",
            "anime": "anime style, manga, vibrant colors, detailed anime art",
            "oil_painting": "oil painting, classical art, brush strokes, artistic masterpiece",
            "watercolor": "watercolor painting, soft colors, artistic, painted texture",
            "cyberpunk": "cyberpunk style, neon lights, futuristic, digital art",
            "fantasy": "fantasy art, magical, ethereal, mystical atmosphere",
            "vintage": "vintage photography, retro, nostalgic, film grain",
            "minimalist": "minimalist art, clean, simple, modern design"
        }
        
        # Prompts negativos padrão
        self.negative_base = "blurry, low quality, distorted, ugly, duplicate, mutation, deformed"
    
    def enhance_prompt(self, base_prompt, style="realism", quality_boost=True):
        """Melhora o prompt com estilo e qualidade"""
        
        # Adicionar estilo
        if style in self.art_styles:
            enhanced = f"{base_prompt}, {self.art_styles[style]}"
        else:
            enhanced = base_prompt
        
        # Adicionar qualidade
        if quality_boost:
            quality_terms = "masterpiece, best quality, highly detailed, sharp focus, professional"
            enhanced = f"{enhanced}, {quality_terms}"
        
        return enhanced
    
    def generate_art(self, prompt, style="realism", width=512, height=512, 
                    num_images=1, steps=20, guidance_scale=7.5, seed=None):
        """Gera arte personalizada"""
        
        print(f"🎨 Gerando arte: '{prompt[:50]}...'")
        print(f"📐 Tamanho: {width}x{height}")
        print(f"🎭 Estilo: {style}")
        
        # Melhorar prompt
        enhanced_prompt = self.enhance_prompt(prompt, style)
        
        # Configurar seed para reproduzibilidade
        generator = None
        if seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        
        # Gerar imagens
        with torch.inference_mode():
            results = self.pipe(
                prompt=enhanced_prompt,
                negative_prompt=self.negative_base,
                width=width,
                height=height,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                num_images_per_prompt=num_images,
                generator=generator
            )
        
        return results.images, enhanced_prompt
    
    def create_art_series(self, theme, variations=5, style="realism"):
        """Cria uma série de arte com variações"""
        
        print(f"🎨 Criando série de arte: {theme}")
        
        # Variações do prompt base
        prompt_variations = [
            f"{theme}, morning light, peaceful atmosphere",
            f"{theme}, dramatic lighting, cinematic view",
            f"{theme}, sunset colors, warm tones",
            f"{theme}, misty atmosphere, ethereal mood",
            f"{theme}, vibrant colors, energetic composition"
        ]
        
        series_results = []
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        series_folder = f"art_series_{timestamp}"
        os.makedirs(series_folder, exist_ok=True)
        
        for i, prompt in enumerate(prompt_variations[:variations], 1):
            print(f"🖼️ Criando variação {i}/{variations}")
            
            images, enhanced_prompt = self.generate_art(
                prompt=prompt,
                style=style,
                seed=42 + i  # Seed diferente para cada variação
            )
            
            # Salvar imagem
            image_filename = f"variation_{i:02d}.png"
            image_path = os.path.join(series_folder, image_filename)
            images[0].save(image_path)
            
            series_results.append({
                'variation': i,
                'prompt': prompt,
                'enhanced_prompt': enhanced_prompt,
                'image_path': image_path
            })
            
            print(f"💾 Salvo: {image_path}")
        
        # Salvar metadata da série
        metadata = {
            'theme': theme,
            'style': style,
            'created_at': datetime.now().isoformat(),
            'variations': series_results
        }
        
        metadata_path = os.path.join(series_folder, "series_info.json")
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)
        
        print(f"🎉 Série completa salva em: {series_folder}")
        return series_results
    
    def create_style_comparison(self, prompt, styles=None):
        """Compara o mesmo prompt em diferentes estilos"""
        
        if styles is None:
            styles = ["realism", "anime", "oil_painting", "cyberpunk"]
        
        print(f"🎭 Comparando estilos para: '{prompt}'")
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        comparison_folder = f"style_comparison_{timestamp}"
        os.makedirs(comparison_folder, exist_ok=True)
        
        comparison_results = []
        
        for style in styles:
            print(f"🎨 Gerando estilo: {style}")
            
            images, enhanced_prompt = self.generate_art(
                prompt=prompt,
                style=style,
                seed=42  # Mesmo seed para comparação justa
            )
            
            # Salvar imagem
            image_filename = f"style_{style}.png"
            image_path = os.path.join(comparison_folder, image_filename)
            images[0].save(image_path)
            
            comparison_results.append({
                'style': style,
                'enhanced_prompt': enhanced_prompt,
                'image_path': image_path
            })
            
            print(f"💾 Salvo: {image_path}")
        
        # Criar grid de comparação
        self.create_comparison_grid(comparison_results, comparison_folder)
        
        return comparison_results
    
    def create_comparison_grid(self, results, folder):
        """Cria um grid com todas as comparações"""
        
        if len(results) < 2:
            return
        
        # Carregar todas as imagens
        images = []
        for result in results:
            img = Image.open(result['image_path'])
            images.append(img)
        
        # Calcular dimensões do grid
        cols = 2
        rows = (len(images) + 1) // 2
        
        img_width, img_height = images[0].size
        grid_width = img_width * cols
        grid_height = img_height * rows
        
        # Criar grid
        grid = Image.new('RGB', (grid_width, grid_height), (255, 255, 255))
        
        for i, img in enumerate(images):
            row = i // cols
            col = i % cols
            x = col * img_width
            y = row * img_height
            grid.paste(img, (x, y))
        
        # Salvar grid
        grid_path = os.path.join(folder, "comparison_grid.png")
        grid.save(grid_path)
        print(f"📊 Grid de comparação salvo: {grid_path}")

# Exemplo de uso completo
if __name__ == "__main__":
    # Inicializar gerador
    generator = ArtisticImageGenerator()
    
    # Exemplo 1: Arte single
    print("\n" + "="*50)
    print("🎨 EXEMPLO 1: ARTE ÚNICA")
    print("="*50)
    
    prompt = "Um dragão majestoso voando sobre um castelo medieval ao pôr do sol"
    images, enhanced = generator.generate_art(
        prompt=prompt,
        style="fantasy",
        width=768,
        height=512
    )
    
    images[0].save("dragao_fantastico.png")
    print(f"✅ Imagem salva: dragao_fantastico.png")
    print(f"📝 Prompt melhorado: {enhanced}")
    
    # Exemplo 2: Série de variações
    print("\n" + "="*50)
    print("🎨 EXEMPLO 2: SÉRIE DE ARTE")
    print("="*50)
    
    series = generator.create_art_series(
        theme="Uma floresta encantada com criaturas mágicas",
        variations=3,
        style="fantasy"
    )
    
    # Exemplo 3: Comparação de estilos
    print("\n" + "="*50)
    print("🎨 EXEMPLO 3: COMPARAÇÃO DE ESTILOS")
    print("="*50)
    
    comparison = generator.create_style_comparison(
        prompt="Um robô futurista em uma cidade cyberpunk",
        styles=["cyberpunk", "anime", "realism", "oil_painting"]
    )
    
    print("\n🎉 TODOS OS EXEMPLOS CONCLUÍDOS!")
    print("📁 Verifique as pastas criadas com suas artes!")