import torch
import torch.nn as nn
from transformers import AutoImageProcessor, AutoModel, MusicgenForConditionalGeneration
import gradio as gr
import numpy as np
import warnings

warnings.filterwarnings("ignore")

VISION_MODEL_ID = "google/siglip-base-patch16-224" 
AUDIO_MODEL_ID = "facebook/musicgen-small"        

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.bfloat16

# =====================================================================
# 🩸 КРОСС-МОДАЛЬНЫЙ МОСТ (Latent-to-Latent)
# =====================================================================
class ChronoLatentBridge(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Загружаем модели с оптимизацией памяти
        self.vision_encoder = AutoModel.from_pretrained(VISION_MODEL_ID, torch_dtype=dtype).to(device)
        self.audio_decoder = MusicgenForConditionalGeneration.from_pretrained(AUDIO_MODEL_ID, torch_dtype=dtype).to(device)
        
        vision_dim = self.vision_encoder.config.vision_config.hidden_size
        audio_dim = self.audio_decoder.config.text_encoder.d_model
        
        # [!] ПАТЧ: Заставляем наш тензорный мост использовать тот же формат (dtype), что и модели
        self.latent_bridge = nn.Linear(vision_dim, audio_dim).to(device, dtype)

    def forward(self, pixel_values):
        vision_outputs = self.vision_encoder.vision_model(pixel_values=pixel_values)
        raw_embeddings = vision_outputs.last_hidden_state 
        
        audio_conditioning = self.latent_bridge(raw_embeddings) 
        seq_len = audio_conditioning.shape[1]
        
        # Monkey Patching (Взлом текстового энкодера)
        original_text_encoder = self.audio_decoder.text_encoder.forward
        
        class VisualThoughts:
            def __init__(self, hidden_states):
                self.last_hidden_state = hidden_states
            def __getitem__(self, idx):
                return [self.last_hidden_state][idx]
                
        def spoofed_text_encoder(*args, **kwargs):
            return VisualThoughts(audio_conditioning)
            
        self.audio_decoder.text_encoder.forward = spoofed_text_encoder
        
        try:
            dummy_inputs = torch.ones((pixel_values.shape[0], seq_len), dtype=torch.long, device=device)
            
            audio_values = self.audio_decoder.generate(
                inputs=dummy_inputs, 
                max_new_tokens=256, 
                do_sample=True,
                guidance_scale=3.5 
            )
        finally:
            self.audio_decoder.text_encoder.forward = original_text_encoder
            
        return audio_values

# =====================================================================

image_processor = AutoImageProcessor.from_pretrained(VISION_MODEL_ID)
chrono_model = ChronoLatentBridge()
sampling_rate = chrono_model.audio_decoder.config.audio_encoder.sampling_rate

def transmute_to_sound(image):
    if image is None: return None
    
    inputs = image_processor(images=image, return_tensors="pt").to(device, dtype)
    
    with torch.no_grad():
        audio_tensor = chrono_model(inputs.pixel_values)
    
    audio_data = audio_tensor[0, 0].cpu().to(torch.float32).numpy()
    audio_data = np.int16(audio_data / np.max(np.abs(audio_data)) * 32767)
    
    return (sampling_rate, audio_data)

PROMO_TEXT = """
**Powered by Livadies. The first artist to synthesize tracks from Cretaceous DNA.**
🟢 [Spotify](https://open.spotify.com/artist/0j8EmbhNFjiVhIJcZHdfUD) | 🔴 [YouTube](https://music.youtube.com/channel/UCe6BJsKd0uj1kAQcdHqyXQw) | 🟡 [Yandex](https://music.yandex.ru/artist/21918652)
🔥 Project Baseline: **«RUSSIAN WINTER 26»**
"""

with gr.Blocks(theme=gr.themes.Monochrome()) as app:
    gr.Markdown("# 🦴 PALEO-SONIC: CHRONO-LATENT ENGINE")
    gr.Markdown("Upload a macro texture of ancient biology (amber, reptile skin, fossils). The custom **Vision-Audio Latent Bridge** bypasses text entirely, translating biological geometry directly into acoustic frequencies.")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="PRIMA MATERIA (Visual Texture)")
            run_btn = gr.Button("TRANSMUTE GEOMETRY TO SOUND", variant="primary")
        with gr.Column():
            out_audio = gr.Audio(label="SYNTHESIZED RESONANCE")
            
    gr.Markdown(PROMO_TEXT)
    run_btn.click(fn=transmute_to_sound, inputs=input_img, outputs=out_audio)

app.launch()