import torch
import torch.nn.functional as F
from torchvision import transforms
import pytesseract # Гипотетически, если бы мы могли это встроить в обучение

def detailed_diffusion_loss(noisy_latents, model_pred, true_noise, text_embeddings, target_text_string, original_image, timesteps):
    """
    Детализированная функция потерь, показывающая, почему текст 'проседает'.
    На практике это сильно упрощено, но концептуально верно.
    """
    # 1. ОСНОВНОЙ DIFFUSION LOSS (Самая большая доля)
    # Минимизирует разницу между предсказанным и реальным шумом.
    # Это ядро обучения диффузионных модель.
    mse_loss = F.mse_loss(model_pred, true_noise)
    # Вес: ~0.55-0.70

    # 2. SEMANTIC ALIGNMENT LOSS (например, через CLIP)
    # Убеждается, что сгенерированное изображение соответствует СМЫСЛУ промпта.
    # Для этого декодируем латентное представление обратно в изображение.
    with torch.no_grad():
        decoded_image = vae.decode(noisy_latents).sample

    # Получаем эмбеддинги изображения и текста через модель CLIP
    clip_image_emb = clip_model.encode_image(transforms.Normalize(...)(decoded_image))
    clip_text_emb = clip_model.encode_text(text_embeddings)
    # Сравниваем их косинусным сходством. Хотим его максимизировать, поэтому используем отрицание.
    clip_loss = -torch.cosine_similarity(clip_image_emb, clip_text_emb).mean()
    # Вес: ~0.20-0.25

    # 3. PERCEPTUAL/VAE RECONSTRUCTION LOSS
    # Отвечает за общее качество изображения, резкость, детализацию.
    # Сравнивает декодированное изображение с оригинальным (до добавления шума) на уровне features.
    perceptual_loss = F.l1_loss(vae.encode(decoded_image).latent_dist.mean,
                                vae.encode(original_image).latent_dist.mean)
    # Вес: ~0.10-0.15

    # 4. TEXT ACCURACY LOSS (ГИПОТЕТИЧЕСКИЙ И ПРОБЛЕМАТИЧНЫЙ)
    # Это тот самый loss, который нам нужен, но его сложно и дорого вычислять.
    text_accuracy_loss = 0.0
    if target_text_string is not None:
        try:
            # Шаг 4.1: Декодируем изображение в пиксельное пространство.
            pil_image = transforms.ToPILImage()(decoded_image.squeeze(0).cpu())
            # Шаг 4.2: Используем внешнюю библиотку OCR (Tesseract) для распознавания текста.
            # Это ОЧЕНЬ медленно и не дифференцируемо по своей природе!
            detected_text = pytesseract.image_to_string(pil_image, config='--psm 8')
            # Шаг 4.3: Вычисляем метрику ошибок (например, Character Error Rate).
            # Это тоже не дифференцируемо.
            cer = calculate_character_error_rate(detected_text, target_text_string)
            text_accuracy_loss = cer
        except Exception as e:
            # OCR может легко упасть, это ненадежный процесс.
            print(f"OCR failed: {e}")
            text_accuracy_loss = 0.0
    # Вес: Вынужденно ставится ~0.01-0.0, потому что он нестабилен и не везде применим.

    # ФАТАЛЬНОЕ ВЗВЕШИВАНИЕ:
    # Именно здесь решается, на что модель будет обращать больше внимания.
    w_mse, w_clip, w_perceptual, w_text = 0.65, 0.22, 0.12, 0.01

    total_loss = (w_mse * mse_loss +
                  w_clip * clip_loss +
                  w_perceptual * perceptual_loss +
                  w_text * text_accuracy_loss)

    return total_loss, {"mse": mse_loss, "clip": clip_loss, "perceptual": perceptual_loss, "text": text_accuracy_loss}

# Вспомогательная функция для CER (недифференцируемая!)
def calculate_character_error_rate(reference, hypothesis):
    # Простейшая реализация CER (расстояние Левенштейна на уровне символов)
    # На практике используются более сложные методы.
    if len(reference) == 0:
        return 1.0 if len(hypothesis) > 0 else 0.0
    # ... (реализация расчета расстояния Левенштейна) ...
    distance = levenshtein_distance(reference, hypothesis)
    return distance / len(reference)

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.attention import CrossAttention

class TextAwareCrossAttention(CrossAttention):
    """
    Модифицированный CrossAttention с механизмом усиления для токенов текста.
    Наследуется от стандартного CrossAttention из библиотеки Diffusers.
    """
    def __init__(self, original_layer, boost_strength=3.0, text_token_ids=None):
        # Инициализируемся параметрами оригинального слоя
        super().__init__(
            query_dim=original_layer.to_q.in_features,
            cross_attention_dim=original_layer.to_k.in_features,
            heads=original_layer.heads,
            dim_head=original_layer.to_q.out_features // original_layer.heads,
            dropout=0.0,
            bias=True,
            upcast_attention=original_layer.upcast_attention,
        )
        
        # Копируем веса из оригинального слоя
        self.load_state_dict(original_layer.state_dict())
        
        self.boost_strength = boost_strength
        self.text_token_ids = text_token_ids if text_token_ids is not None else []
        
    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
        """
        Args:
            hidden_states: [batch_size, sequence_length, channels] - латентное представление изображения
            encoder_hidden_states: [batch_size, text_seq_len, cross_attention_dim] - текстовые эмбеддинги
        """
        # Стандартный forward до расчета внимания
        batch_size, sequence_length, _ = hidden_states.shape
        
        query = self.to_q(hidden_states)
        key = self.to_k(encoder_hidden_states)
        value = self.to_v(encoder_hidden_states)

        # Перестраиваем для multi-head attention
        query = self.reshape_heads_to_batch_dim(query)
        key = self.reshape_heads_to_batch_dim(key)
        value = self.reshape_heads_to_batch_dim(value)

        # 1. Расчет матрицы внимания
        attention_scores = torch.baddbmm(
            torch.empty(query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device),
            query,
            key.transpose(-1, -2),
            beta=0,
            alpha=self.scale,
        )
        
        # 2. КРИТИЧЕСКИЙ ЭТАП: Создание маски усиления для текстовых токенов
        if self.text_token_ids and encoder_hidden_states is not None:
            text_seq_len = encoder_hidden_states.shape[1]
            boost_mask = torch.zeros_like(attention_scores)
            
            # Создаем маску, где текстовые токены получают усиление
            for token_id in self.text_token_ids:
                if token_id < text_seq_len:
                    # Усиливаем внимание ко ВСЕМ текстовым токенам
                    boost_mask[:, :, token_id] = self.boost_strength
            
            # Применяем маску: усиливаем scores для целевых токенов
            attention_scores = attention_scores + boost_mask

        # 3. Применяем softmax к модифицированным scores
        attention_probs = F.softmax(attention_scores, dim=-1)
        
        # 4. Стандартный расчет выходных значений
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
        hidden_states = self.to_out[0](hidden_states)
        hidden_states = self.to_out[1](hidden_states)
        
        return hidden_states

def inject_text_aware_attention(pipeline, target_text, tokenizer):
    """
    Функция для внедрения нашего контролируемого внимания в pipeline.
    """
    # 1. Токенизируем целевой текст, чтобы найти индексы токенов
    tokens = tokenizer(
        target_text,
        padding="do_not_padding",
        truncation=True,
        return_tensors="pt",
    )
    text_token_ids = tokens.input_ids[0].tolist()
    
    # 2. Рекурсивно обходим U-Net и заменяем cross-attention слои
    def replace_attention_layers(module, text_token_ids):
        for name, child in module.named_children():
            if isinstance(child, CrossAttention) and child.cross_attention_dim is not None:
                # Заменяем стандартный слой на наш кастомный
                new_layer = TextAwareCrossAttention(child, boost_strength=4.0, text_token_ids=text_token_ids)
                setattr(module, name, new_layer)
            else:
                # Рекурсивно применяем к дочерним модулям
                replace_attention_layers(child, text_token_ids)
    
    replace_attention_layers(pipeline.unet, text_token_ids)
    return pipeline

from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Внедряем контроль внимания для текста "PHOENIX CAFE"
pipe = inject_text_aware_attention(pipe, "PHOENIX CAFE", pipe.tokenizer)

# Генерируем изображение с усиленным вниманием к тексту
prompt = "a vintage sign that says 'PHOENIX CAFE' on a brick wall"
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

def create_text_mask(width, height, text, font_size=60, font_path="arial.ttf"):
    """Создает белую маску с черным текстом для ControlNet."""
    # Создаем черное изображение
    mask = Image.new("L", (width, height), 0)
    draw = ImageDraw.Draw(mask)
    
    try:
        font = ImageFont.truetype(font_path, font_size)
    except:
        font = ImageFont.load_default()
    
    # Получаем bounding box текста
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    # Вычисляем позицию для центрирования
    x = (width - text_width) / 2
    y = (height - text_height) / 2
    
    # Рисуем БЕЛЫЙ текст на ЧЕРНОМ фоне
    draw.text((x, y), text, fill=255, font=font)
    
    return mask

def create_scribble_mask(width, height, text, thickness=2):
    """Создает маску в стиле скетча/наброска."""
    # Сначала создаем обычную текстовую маску
    text_mask = create_text_mask(width, height, text)
    
    # Конвертируем в numpy для OpenCV обработки
    mask_np = np.array(text_mask)
    
    # Находим контуры текста
    contours, _ = cv2.findContours(mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Создаем чистую маску и рисуем только контуры
    scribble_mask = np.zeros_like(mask_np)
    cv2.drawContours(scribble_mask, contours, -1, 255, thickness)
    
    return Image.fromarray(scribble_mask)

# ЗАГРУЗКА И НАСТРОЙКА ПАЙПЛАЙНА
controlnet1 = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0",
    torch_dtype=torch.float16
)
controlnet2 = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-scribble",
    torch_dtype=torch.float16
)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=[controlnet1, controlnet2],
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

# ПОДГОТОВКА МАСОК
width, height = 1024, 1024
target_text = "PHOENIX\nCOFFEE"

# Маска 1: Canny edges для сохранения структуры
text_mask = create_text_mask(width, height, target_text)
canny_mask = cv2.Canny(np.array(text_mask), 100, 200)
canny_mask = Image.fromarray(canny_mask)

# Маска 2: Scribble для стилистического руководства
scribble_mask = create_scribble_mask(width, height, target_text, thickness=4)

# ГЕНЕРАЦИЯ С ДВУМЯ CONTROLNET
prompt = "a beautiful vintage coffee shop sign, high quality, detailed, 'PHOENIX COFFEE' text, gold letters, black background"
negative_prompt = "blurry, low quality, distorted text, bad typography"

image = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=[canny_mask, scribble_mask],  # Две маски для двух ControlNet
    num_inference_steps=30,
    guidance_scale=8.0,
    controlnet_conditioning_scale=[0.7, 0.5],  # Разные веса для разных масок
).images[0]

import os
import json
import random
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageOps
import numpy as np
from pathlib import Path

class AdvancedTextDatasetGenerator:
    def __init__(self, output_dir="synthetic_dataset", fonts_dir="fonts"):
        self.output_dir = Path(output_dir)
        self.fonts_dir = Path(fonts_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # Загружаем все доступные шрифты
        self.font_paths = list(self.fonts_dir.glob("*.ttf")) + list(self.fonts_dir.glob("*.otf"))
        if not self.font_paths:
            raise ValueError(f"No fonts found in {fonts_dir}")
        
        # База слов для осмысленных текстов
        self.meaningful_words = ["CAFE", "BAR", "RESTAURANT", "HOTEL", "PHOENIX", "DRAGON", "ROYAL", "GRAND", "CENTRAL", "URBAN"]
        
        # Случайные последовательности для обобщения
        self.random_chars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
        
    def create_complex_background(self, width, height):
        """Создает сложный фон с градиентами, текстурами и шумом."""
        # Вариант 1: Градиентный фон
        if random.random() < 0.3:
            bg = Image.new('RGB', (width, height))
            draw = ImageDraw.Draw(bg)
            
            # Случайный градиент
            for i in range(height):
                ratio = i / height
                r = int(random.randint(0, 100) * (1 - ratio) + random.randint(150, 255) * ratio)
                g = int(random.randint(0, 100) * (1 - ratio) + random.randint(150, 255) * ratio)
                b = int(random.randint(0, 100) * (1 - ratio) + random.randint(150, 255) * ratio)
                draw.line([(0, i), (width, i)], fill=(r, g, b))
        
        # Вариант 2: Текстурный фон (дерево, металл, камень)
        elif random.random() < 0.5:
            # Создаем базовый шум и применяем фильтры для имитации текстуры
            bg = Image.new('RGB', (width, height))
            pixels = np.random.randint(50, 150, (height, width, 3), dtype=np.uint8)
            bg = Image.fromarray(pixels)
            
            # Применяем размытие и шум для создания текстуры
            if random.random() < 0.5:
                bg = bg.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.5, 2.0)))
            
            # Добавляем шум
            noise = np.random.randint(0, 30, (height, width, 3), dtype=np.uint8)
            bg = Image.fromarray(np.clip(np.array(bg) + noise, 0, 255).astype(np.uint8))
        
        # Вариант 3: Простой цветной фон
        else:
            bg_color = (random.randint(200, 255), random.randint(200, 255), random.randint(200, 255))
            bg = Image.new('RGB', (width, height), bg_color)
            
        return bg
    
    def add_text_effects(self, draw, text, font, x, y, fill_color):
        """Добавляет визуальные эффекты к тексту."""
        # Эффект тени
        if random.random() < 0.4:
            shadow_color = (0, 0, 0) if random.random() < 0.5 else (50, 50, 50)
            shadow_offset = random.randint(2, 4)
            draw.text((x + shadow_offset, y + shadow_offset), text, font=font, fill=shadow_color)
        
        # Эффект обводки
        if random.random() < 0.3:
            stroke_width = random.randint(1, 3)
            stroke_color = (0, 0, 0) if max(fill_color) > 128 else (255, 255, 255)
            
            # Рисуем обводку в нескольких направлениях
            for dx in [-stroke_width, 0, stroke_width]:
                for dy in [-stroke_width, 0, stroke_width]:
                    if dx != 0 or dy != 0:
                        draw.text((x + dx, y + dy), text, font=font, fill=stroke_color)
        
        # Основной текст
        draw.text((x, y), text, font=font, fill=fill_color)
    
    def generate_sample(self, sample_id, width=1024, height=1024):
        """Генерирует один sample синтетических данных."""
        # 1. Создаем сложный фон
        image = self.create_complex_background(width, height)
        draw = ImageDraw.Draw(image)
        
        # 2. Выбираем тип текста: осмысленный или случайный
        if random.random() < 0.7:
            # Осмысленный текст (1-3 слова)
            num_words = random.randint(1, 3)
            text = ' '.join(random.sample(self.meaningful_words, num_words))
        else:
            # Случайная последовательность
            text_length = random.randint(3, 8)
            text = ''.join(random.choices(self.random_chars, k=text_length))
        
        # 3. Выбираем шрифт и размер
        font_path = random.choice(self.font_paths)
        font_size = random.randint(40, 120)
        
        try:
            font = ImageFont.truetype(str(font_path), font_size)
        except:
            # Fallback шрифт
            font = ImageFont.load_default()
        
        # 4. Рассчитываем позиционирование
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        
        # Случайная позиция с отступами от краев
        margin_x = random.randint(50, 200)
        margin_y = random.randint(50, 200)
        x = random.randint(margin_x, width - text_width - margin_x)
        y = random.randint(margin_y, height - text_height - margin_y)
        
        # 5. Выбираем цвет текста (контрастный к фону)
        bg_color = image.getpixel((x, y))
        # Убеждаемся в контрастности
        if sum(bg_color) > 384:  # Светлый фон
            text_color = (random.randint(0, 100), random.randint(0, 100), random.randint(0, 100))
        else:  # Темный фон
            text_color = (random.randint(155, 255), random.randint(155, 255), random.randint(155, 255))
        
        # 6. Добавляем текст с эффектами
        self.add_text_effects(draw, text, font, x, y, text_color)
        
        # 7. Иногда добавляем дополнительные эффекты ко всему изображению
        if random.random() < 0.2:
            image = image.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.1, 0.5)))
        
        # 8. Сохраняем изображение
        img_filename = f"sample_{sample_id:06d}.png"
        image.save(self.output_dir / img_filename, "PNG")
        
        # 9. Создаем метаданные
        metadata = {
            "file_name": img_filename,
            "text": text,
            "font": font_path.name,
            "font_size": font_size,
            "text_color": text_color,
            "position": {"x": int(x), "y": int(y)},
            "dimensions": {"width": text_width, "height": text_height},
            "background_type": "complex"
        }
        
        return metadata
    
    def generate_dataset(self, num_samples=10000):
        """Генерирует полный датасет."""
        metadata_list = []
        
        for i in range(num_samples):
            if i % 1000 == 0:
                print(f"Generated {i}/{num_samples} samples...")
            
            try:
                metadata = self.generate_sample(i)
                metadata_list.append(metadata)
            except Exception as e:
                print(f"Error generating sample {i}: {e}")
                continue
        
        # Сохраняем метаданные
        with open(self.output_dir / "metadata.jsonl", "w") as f:
            for meta in metadata_list:
                f.write(json.dumps(meta) + "\n")
        
        print(f"Dataset generation complete. {len(metadata_list)} samples created.")

# ИСПОЛЬЗОВАНИЕ
generator = AdvancedTextDatasetGenerator(
    output_dir="my_synthetic_text_dataset",
    fonts_dir="path/to/your/fonts"  # Папка с .ttf/.otf файлами
)
generator.generate_dataset(num_samples=50000)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from transformers import CLIPModel, CLIPProcessor
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2

class ComprehensiveTextAccuracyLoss(nn.Module):
    """
    Полностью реализованная кастомная функция потерь для улучшения читаемости текста.
    Комбинирует дифференцируемые подходы для обхода проблемы недифференцируемости OCR.
    """
    def __init__(self, clip_model_name="openai/clip-vit-base-patch32", device="cuda"):
        super().__init__()
        self.device = device
        self.clip_model = CLIPModel.from_pretrained(clip_model_name).to(device)
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
        
        # Замораживаем CLIP
        for param in self.clip_model.parameters():
            param.requires_grad = False
            
        # Инициализируем дифференцируемый рендерер текста
        self.text_renderer = DifferentiableTextRenderer(device=device)
        
        # Настраиваем веса для разных компонентов loss
        self.weights = {
            'clip_consistency': 0.3,
            'text_aware_clip': 0.3,
            'structural_similarity': 0.2,
            'edge_consistency': 0.2
        }
            
    def forward(self, generated_images, target_texts, original_prompts):
        batch_size = generated_images.shape[0]
        total_loss = torch.tensor(0.0, device=self.device)
        
        for i in range(batch_size):
            # 1. CLIP Text-Image Consistency Loss
            clip_loss = self.compute_clip_consistency(
                generated_images[i].unsqueeze(0), 
                target_texts[i]
            )
            
            # 2. Text-Aware CLIP Loss
            text_aware_loss = self.compute_text_aware_clip(
                generated_images[i].unsqueeze(0),
                original_prompts[i],
                target_texts[i]
            )
            
            # 3. Structural Similarity Loss
            structural_loss = self.compute_structural_similarity(
                generated_images[i].unsqueeze(0),
                target_texts[i]
            )
            
            # 4. Edge Consistency Loss
            edge_loss = self.compute_edge_consistency(
                generated_images[i].unsqueeze(0),
                target_texts[i]
            )
            
            # Комбинируем все компоненты с весами
            sample_loss = (
                self.weights['clip_consistency'] * clip_loss +
                self.weights['text_aware_clip'] * text_aware_loss +
                self.weights['structural_similarity'] * structural_loss +
                self.weights['edge_consistency'] * edge_loss
            )
            
            total_loss += sample_loss
        
        return total_loss / batch_size
    
    def compute_clip_consistency(self, image, target_text):
        """Loss на основе CLIP: насколько изображение соответствует целевому тексту."""
        inputs = self.clip_processor(
            text=[target_text], 
            images=image, 
            return_tensors="pt", 
            padding=True
        ).to(self.device)
        
        outputs = self.clip_model(**inputs)
        similarity = F.cosine_similarity(outputs.image_embeds, outputs.text_embeds)
        return 1 - similarity.mean()
    
    def compute_text_aware_clip(self, image, original_prompt, target_text):
        """Loss, который усиливает важность текстовой части промпта."""
        enhanced_prompt = f"{original_prompt} with clear, readable text that says '{target_text}'"
        
        inputs_normal = self.clip_processor(
            text=[original_prompt], 
            images=image, 
            return_tensors="pt", 
            padding=True
        ).to(self.device)
        
        inputs_enhanced = self.clip_processor(
            text=[enhanced_prompt], 
            images=image, 
            return_tensors="pt", 
            padding=True
        ).to(self.device)
        
        outputs_normal = self.clip_model(**inputs_normal)
        outputs_enhanced = self.clip_model(**inputs_enhanced)
        
        sim_normal = F.cosine_similarity(outputs_normal.image_embeds, outputs_normal.text_embeds)
        sim_enhanced = F.cosine_similarity(outputs_enhanced.image_embeds, outputs_enhanced.text_embeds)
        
        return F.relu(sim_normal - sim_enhanced + 0.1)
    
    def compute_structural_similarity(self, image, target_text):
        """Loss на основе структурного сходства с идеально отрендеренным текстом."""
        # Рендерим идеальный текст с теми же размерами
        ideal_text_image = self.text_renderer.render_text_batch(
            [target_text], 
            image.shape[2],  # height
            image.shape[3]   # width
        )
        
        # Вычисляем структурное сходство (SSIM)
        ssim_loss = 1 - self.ssim(image, ideal_text_image)
        return ssim_loss
    
    def compute_edge_consistency(self, image, target_text):
        """Loss на основе согласованности границ текста."""
        # Рендерим идеальный текст для сравнения границ
        ideal_text_image = self.text_renderer.render_text_batch(
            [target_text], 
            image.shape[2], 
            image.shape[3]
        )
        
        # Вычисляем границы с помощью дифференцируемого оператора Собеля
        generated_edges = self.sobel_edges(image)
        ideal_edges = self.sobel_edges(ideal_text_image)
        
        # Сравниваем границы с помощью MSE
        edge_loss = F.mse_loss(generated_edges, ideal_edges)
        return edge_loss
    
    def ssim(self, x, y, window_size=11, size_average=True):
        """Вычисляет Structural Similarity Index (SSIM)."""
        from math import exp
        
        # Параметры SSIM
        C1 = 0.01 ** 2
        C2 = 0.03 ** 2
        
        mu_x = F.avg_pool2d(x, window_size, stride=1, padding=window_size//2)
        mu_y = F.avg_pool2d(y, window_size, stride=1, padding=window_size//2)
        
        mu_x_sq = mu_x.pow(2)
        mu_y_sq = mu_y.pow(2)
        mu_x_mu_y = mu_x * mu_y
        
        sigma_x_sq = F.avg_pool2d(x * x, window_size, stride=1, padding=window_size//2) - mu_x_sq
        sigma_y_sq = F.avg_pool2d(y * y, window_size, stride=1, padding=window_size//2) - mu_y_sq
        sigma_xy = F.avg_pool2d(x * y, window_size, stride=1, padding=window_size//2) - mu_x_mu_y
        
        ssim_n = (2 * mu_x_mu_y + C1) * (2 * sigma_xy + C2)
        ssim_d = (mu_x_sq + mu_y_sq + C1) * (sigma_x_sq + sigma_y_sq + C2)
        ssim = ssim_n / ssim_d
        
        return ssim.mean() if size_average else ssim
    
    def sobel_edges(self, x):
        """Вычисляет границы с помощью оператора Собеля."""
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=self.device).view(1, 1, 3, 3)
        sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=self.device).view(1, 1, 3, 3)
        
        # Применяем к каждому каналу
        edges_x = torch.zeros_like(x)
        edges_y = torch.zeros_like(x)
        
        for i in range(x.shape[1]):
            edges_x[:, i:i+1] = F.conv2d(x[:, i:i+1], sobel_x, padding=1)
            edges_y[:, i:i+1] = F.conv2d(x[:, i:i+1], sobel_y, padding=1)
        
        # Объединяем границы
        edges = torch.sqrt(edges_x ** 2 + edges_y ** 2)
        return edges

class DifferentiableTextRenderer(nn.Module):
    """Дифференцируемый рендерер текста для использования в функциях потерь."""
    
    def __init__(self, device="cuda"):
        super().__init__()
        self.device = device
        
        # Создаем базовые шрифты разных размеров
        self.font_sizes = [24, 36, 48, 64]
        self.fonts = []
        
        for size in self.font_sizes:
            try:
                # Пытаемся загрузить шрифт (нужно иметь .ttf файлы в системе)
                font = ImageFont.truetype("arial.ttf", size)
                self.fonts.append(font)
            except:
                # Fallback на стандартный шрифт
                font = ImageFont.load_default()
                self.fonts.append(font)
    
    def render_text_batch(self, texts, height, width):
        """Рендерит батч текстов в тензоры."""
        batch_size = len(texts)
        rendered_batch = torch.zeros(batch_size, 3, height, width, device=self.device)
        
        for i, text in enumerate(texts):
            # Рендерим каждый текст отдельно
            text_tensor = self.render_single_text(text, height, width)
            rendered_batch[i] = text_tensor
        
        return rendered_batch
    
    def render_single_text(self, text, height, width):
        """Рендерит один текст в тензор."""
        # Создаем PIL изображение
        pil_image = Image.new('RGB', (width, height), color=(255, 255, 255))
        draw = ImageDraw.Draw(pil_image)
        
        # Выбираем случайный шрифт
        font = np.random.choice(self.fonts)
        
        # Получаем bounding box текста
        bbox = draw.textbbox((0, 0), text, font=font)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]
        
        # Центрируем текст
        x = (width - text_width) / 2
        y = (height - text_height) / 2
        
        # Рисуем черный текст на белом фоне
        draw.text((x, y), text, fill=(0, 0, 0), font=font)
        
        # Конвертируем в тензор
        image_tensor = transforms.ToTensor()(pil_image).to(self.device)
        return image_tensor

# ПОЛНАЯ ИМПЛЕМЕНТАЦИЯ ПРОЦЕССА ОБУЧЕНИЯ
def train_with_text_enhancement(model, train_dataloader, val_dataloader, num_epochs=10):
    """Полная функция обучения с улучшением генерации текста."""
    
    # Инициализируем наши кастомные компоненты
    text_loss_fn = ComprehensiveTextAccuracyLoss(device="cuda")
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
    
    # Метрики
    train_losses = []
    val_losses = []
    text_accuracy_metrics = []
    
    for epoch in range(num_epochs):
        print(f"Epoch {epoch+1}/{num_epochs}")
        
        # Фаза обучения
        model.train()
        epoch_train_loss = 0
        epoch_text_loss = 0
        
        for batch_idx, batch in enumerate(train_dataloader):
            optimizer.zero_grad()
            
            # Подготавливаем данные
            latents = batch["latents"].to("cuda")
            noise = batch["noise"].to("cuda") 
            timesteps = batch["timesteps"].to("cuda")
            text_embeddings = batch["text_embeddings"].to("cuda")
            target_texts = batch["target_texts"]
            prompts = batch["prompts"]
            
            # Forward pass модели
            noise_pred = model(latents, timesteps, text_embeddings).sample
            
            # 1. Стандартный diffusion loss
            mse_loss = F.mse_loss(noise_pred, noise)
            
            # 2. Наш кастомный text accuracy loss
            with torch.no_grad():
                # Декодируем латенты в изображения для text loss
                generated_images = decode_latents_to_pixels(latents)
            
            text_loss = text_loss_fn(generated_images, target_texts, prompts)
            
            # Комбинируем losses
            total_loss = 0.7 * mse_loss + 0.3 * text_loss
            
            # Backward pass
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            
            epoch_train_loss += total_loss.item()
            epoch_text_loss += text_loss.item()
            
            if batch_idx % 100 == 0:
                print(f"Batch {batch_idx}, Total Loss: {total_loss.item():.4f}, Text Loss: {text_loss.item():.4f}")
        
        # Фаза валидации
        model.eval()
        epoch_val_loss = 0
        val_text_accuracy = 0
        
        with torch.no_grad():
            for val_batch in val_dataloader:
                val_latents = val_batch["latents"].to("cuda")
                val_noise = val_batch["noise"].to("cuda")
                val_timesteps = val_batch["timesteps"].to("cuda")
                val_text_embeddings = val_batch["text_embeddings"].to("cuda")
                val_target_texts = val_batch["target_texts"]
                val_prompts = val_batch["prompts"]
                
                val_noise_pred = model(val_latents, val_timesteps, val_text_embeddings).sample
                val_mse_loss = F.mse_loss(val_noise_pred, val_noise)
                
                val_generated_images = decode_latents_to_pixels(val_latents)
                val_text_loss = text_loss_fn(val_generated_images, val_target_texts, val_prompts)
                
                val_total_loss = 0.7 * val_mse_loss + 0.3 * val_text_loss
                epoch_val_loss += val_total_loss.item()
                
                # Вычисляем accuracy текста (используя OCR для валидации)
                text_accuracy = compute_text_accuracy_ocr(val_generated_images, val_target_texts)
                val_text_accuracy += text_accuracy
        
        # Вычисляем средние метрики эпохи
        avg_train_loss = epoch_train_loss / len(train_dataloader)
        avg_val_loss = epoch_val_loss / len(val_dataloader)
        avg_text_accuracy = val_text_accuracy / len(val_dataloader)
        
        train_losses.append(avg_train_loss)
        val_losses.append(avg_val_loss)
        text_accuracy_metrics.append(avg_text_accuracy)
        
        print(f"Epoch {epoch+1} Summary:")
        print(f"Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
        print(f"Text Accuracy: {avg_text_accuracy:.4f}")
        
        # Сохраняем чекпоинт
        if (epoch + 1) % 5 == 0:
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'train_loss': avg_train_loss,
                'val_loss': avg_val_loss,
                'text_accuracy': avg_text_accuracy
            }
            torch.save(checkpoint, f'text_enhanced_model_epoch_{epoch+1}.pth')
        
        scheduler.step()
    
    return {
        'train_losses': train_losses,
        'val_losses': val_losses,
        'text_accuracy': text_accuracy_metrics,
        'final_model': model
    }

def decode_latents_to_pixels(latents):
    """Декодирует латенты обратно в пиксельное пространство."""
    # Эта функция зависит от конкретной реализации VAE в вашей диффузионной модели
    # Здесь приведен упрощенный пример
    scale_factor = 0.18215  # Стандартный scale factor для Stable Diffusion
    latents = latents / scale_factor
    
    # Используем VAE для декодирования
    with torch.no_grad():
        images = vae.decode(latents).sample
    
    # Нормализуем изображения в [0, 1]
    images = (images / 2 + 0.5).clamp(0, 1)
    return images

def compute_text_accuracy_ocr(generated_images, target_texts):
    """Вычисляет accuracy текста с помощью OCR (только для валидации)."""
    total_accuracy = 0
    batch_size = generated_images.shape[0]
    
    for i in range(batch_size):
        # Конвертируем тензор в PIL Image
        image_tensor = generated_images[i].cpu()
        image_pil = transforms.ToPILImage()(image_tensor)
        
        try:
            # Используем pytesseract для OCR
            import pytesseract
            detected_text = pytesseract.image_to_string(image_pil, config='--psm 8')
            
            # Простое сравнение текстов
            target = target_texts[i].upper().strip()
            detected = detected_text.upper().strip()
            
            if target in detected or detected in target:
                total_accuracy += 1
        except:
            # Если OCR не работает, пропускаем этот sample
            continue
    
    return total_accuracy / batch_size

# ИНИЦИАЛИЗАЦИЯ И ЗАПУСК ОБУЧЕНИЯ
def main():
    """Основная функция для запуска процесса обучения."""
    
    # Загружаем предобученную модель
    from diffusers import StableDiffusionPipeline
    model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    
    # Подготавливаем датасет
    train_dataset = TextEnhancedDataset("synthetic_dataset/train")
    val_dataset = TextEnhancedDataset("synthetic_dataset/val")
    
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=4, shuffle=True)
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=4, shuffle=False)
    
    # Запускаем обучение
    results = train_with_text_enhancement(
        model=model.unet,  # Обучаем только U-Net
        train_dataloader=train_loader,
        val_dataloader=val_loader,
        num_epochs=20
    )
    
    print("Training completed!")
    print(f"Final text accuracy: {results['text_accuracy'][-1]:.4f}")

if __name__ == "__main__":
    main()

class ComprehensiveTextGenerationPipeline:
    """
    Комплексный пайплайн, объединяющий все методы для максимального качества текста.
    """
    
    def __init__(self, model_name="runwayml/stable-diffusion-v1-5", device="cuda"):
        self.device = device
        
        # Загружаем базовую модель
        self.base_pipeline = StableDiffusionPipeline.from_pretrained(
            model_name, 
            torch_dtype=torch.float16
        ).to(device)
        
        # Загружаем дообученную модель (если есть)
        try:
            self.fine_tuned_model = self.load_fine_tuned_model()
            self.use_fine_tuned = True
        except:
            self.use_fine_tuned = False
            print("Fine-tuned model not found, using base model")
        
        # Инициализируем контроллеры внимания
        self.attention_controller = AttentionController()
        
        # Загружаем ControlNet модели
        self.controlnet_models = self.load_controlnet_models()
        
    def generate_with_text_control(self, prompt, target_text, 
                                 use_attention_control=True,
                                 use_controlnet=True,
                                 controlnet_strength=0.7,
                                 num_inference_steps=50,
                                 guidance_scale=7.5):
        """
        Генерирует изображение с полным контролем над текстом.
        """
        
        # 1. Подготовка пайплайна
        if self.use_fine_tuned:
            pipeline = self.fine_tuned_model
        else:
            pipeline = self.base_pipeline
        
        # 2. Применяем контроль внимания
        if use_attention_control:
            pipeline = self.attention_controller.inject_text_attention(
                pipeline, target_text, pipeline.tokenizer
            )
        
        # 3. Подготавливаем ControlNet маски
        controlnet_images = []
        controlnet_models = []
        
        if use_controlnet and self.controlnet_models:
            # Создаем текстовую маску
            text_mask = self.create_advanced_text_mask(512, 512, target_text)
            
            # Добавляем разные типы ControlNet для лучшего контроля
            canny_mask = self.create_canny_mask(text_mask)
            scribble_mask = self.create_scribble_mask(text_mask)
            
            controlnet_images.extend([canny_mask, scribble_mask])
            controlnet_models.extend([
                self.controlnet_models['canny'],
                self.controlnet_models['scribble']
            ])
        
        # 4. Генерация
        if controlnet_models:
            # Генерация с ControlNet
            from diffusers import StableDiffusionControlNetPipeline
            
            controlnet_pipeline = StableDiffusionControlNetPipeline(
                vae=pipeline.vae,
                text_encoder=pipeline.text_encoder,
                tokenizer=pipeline.tokenizer,
                unet=pipeline.unet,
                scheduler=pipeline.scheduler,
                safety_checker=pipeline.safety_checker,
                feature_extractor=pipeline.feature_extractor,
                controlnet=controlnet_models,
            ).to(self.device)
            
            image = controlnet_pipeline(
                prompt=prompt,
                image=controlnet_images,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                controlnet_conditioning_scale=[controlnet_strength] * len(controlnet_models),
                height=512,
                width=512,
            ).images[0]
        else:
            # Стандартная генерация
            image = pipeline(
                prompt=prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                height=512,
                width=512,
            ).images[0]
        
        return image
    
    def create_advanced_text_mask(self, width, height, text):
        """Создает продвинутую текстовую маску с разными эффектами."""
        # Реализация создания сложной маски...
        pass
    
    def load_fine_tuned_model(self):
        """Загружает дообученную модель."""
        # Реализация загрузки модели...
        pass
    
    def load_controlnet_models(self):
        """Загружает различные ControlNet модели."""
        controlnets = {}
        try:
            controlnets['canny'] = ControlNetModel.from_pretrained(
                "lllyasviel/sd-controlnet-canny"
            ).to(self.device)
            controlnets['scribble'] = ControlNetModel.from_pretrained(
                "lllyasviel/sd-controlnet-scribble" 
            ).to(self.device)
        except Exception as e:
            print(f"Error loading ControlNet models: {e}")
        
        return controlnets

# ПРИМЕР ИСПОЛЬЗОВАНИЯ КОМПЛЕКСНОГО ПАЙПЛАЙНА
def demonstrate_comprehensive_pipeline():
    """Демонстрирует работу комплексного пайплайна."""
    
    pipeline = ComprehensiveTextGenerationPipeline()
    
    # Разные сценарии генерации
    scenarios = [
        {
            'prompt': 'a vintage coffee shop sign on a brick wall',
            'target_text': 'PHOENIX CAFE',
            'use_attention_control': True,
            'use_controlnet': True
        },
        {
            'prompt': 'modern tech company logo, clean design',
            'target_text': 'NEXUS AI',
            'use_attention_control': True,
            'use_controlnet': False
        },
        {
            'prompt': 'restaurant menu board, chalkboard style',
            'target_text': 'SPECIALS\nPASTA $12\nSALAD $8',
            'use_attention_control': True,
            'use_controlnet': True
        }
    ]
    
    for i, scenario in enumerate(scenarios):
        print(f"Generating image {i+1}/{len(scenarios)}...")
        
        image = pipeline.generate_with_text_control(
            prompt=scenario['prompt'],
            target_text=scenario['target_text'],
            use_attention_control=scenario['use_attention_control'],
            use_controlnet=scenario['use_controlnet']
        )
        
        # Сохраняем результат
        image.save(f"comprehensive_result_{i+1}.png")
        
        # Оцениваем качество текста
        accuracy = evaluate_text_quality(image, scenario['target_text'])
        print(f"Text accuracy for image {i+1}: {accuracy:.4f}")

def evaluate_text_quality(image, target_text):
    """Оценивает качество сгенерированного текста."""
    try:
        import pytesseract
        
        # Извлекаем текст с изображения
        detected_text = pytesseract.image_to_string(image, config='--psm 8')
        
        # Простое сравнение (можно улучшить)
        target_clean = target_text.upper().replace('\n', ' ').strip()
        detected_clean = detected_text.upper().replace('\n', ' ').strip()
        
        if target_clean in detected_clean:
            return 1.0
        else:
            # Вычисляем схожесть
            from difflib import SequenceMatcher
            similarity = SequenceMatcher(None, target_clean, detected_clean).ratio()
            return similarity
            
    except Exception as e:
        print(f"Error in text evaluation: {e}")
        return 0.0

# Запуск демонстрации
if __name__ == "__main__":
    demonstrate_comprehensive_pipeline()