Update app.py

app.py

@@ -8,12 +8,10 @@ import matplotlib.animation as animation
 import time
 import threading
 from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModel, TrainingArguments, pipeline
+from transformers import AutoTokenizer, AutoModel, AutoModelForTextToWaveform, TrainingArguments, pipeline
 from diffusers import DiffusionPipeline
 from huggingface_hub import login, HfApi, Repository
 from dotenv import load_dotenv
-import gradio as gr
-
 
 # Cargar variables de entorno
 load_dotenv()
@@ -114,6 +112,9 @@ def main():
 
             # Inicializar los pipelines
             pipelines_to_unify = [
+                pipeline("text-to-audio", model="facebook/musicgen-melody"),
+                pipeline("text-to-audio", model="facebook/musicgen-large"),
+                pipeline("text-to-audio", model="facebook/musicgen-small"),
                 DiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt-1-1"),
                 pipeline("automatic-speech-recognition", model="openai/whisper-small"),
                 DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev"),
@@ -134,6 +135,13 @@ def main():
                 pipeline("text-generation", model="Groq/Llama-3-Groq-8B-Tool-Use")
             ]
 
+            # Añadir modelos adicionales
+            additional_models = [
+                "facebook/musicgen-large",
+                "facebook/musicgen-melody"
+            ]
+
+            # Inicializar los tokenizadores y modelos
             tokenizers = {}
             models = []
             for model_name in models_to_train:
@@ -146,8 +154,11 @@ def main():
                 tokenizers[model_name] = tokenizer
                 models.append(model)
 
-            # Agregar pipelines como modelos
-            models.extend(pipelines_to_unify)
+            for model_name in additional_models:
+                tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+                model = AutoModelForTextToWaveform.from_pretrained(model_name)
+                tokenizers[model_name] = tokenizer
+                models.append(model)
 
             # Crear un dataset sintético para entrenamiento y evaluación
             synthetic_dataset = SyntheticDataset(tokenizers, size=100)
@@ -191,45 +202,43 @@ def main():
             def train(model, train_loader, eval_loader, args):
                 model.train()
                 epoch = 0
-                total_steps = args.num_train_epochs * len(train_loader)
-                progress_bar = tqdm(total=total_steps, desc="Training")
-
-                while epoch < args.num_train_epochs:
+                total_steps = len(train_loader)
+                for step, batch in enumerate(train_loader):
                     start_time = time.time()
-                    for step, batch in enumerate(train_loader):
+                    input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
+                    attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
+                    labels = batch["label"].to("cpu")
+                    optimizer.zero_grad()
+                    outputs = model(input_ids)
+                    loss = nn.CrossEntropyLoss()(outputs, labels)
+                    loss.backward()
+                    optimizer.step()
+
+                    progress_bar.update(1)
+
+                    elapsed_time = time.time() - start_time
+                    estimated_total_time = total_steps * (elapsed_time / (step + 1))
+                    estimated_remaining_time = estimated_total_time - elapsed_time
+
+                    if step % args.logging_steps == 0:
+                        train_losses.append(loss.item())
+                        print(f"Step {step}/{total_steps}, Loss: {loss.item()}, Estimated remaining time: {estimated_remaining_time:.2f} seconds")
+
+                epoch += 1
+                model.eval()
+                eval_loss = 0
+                with torch.no_grad():
+                    for batch in eval_loader:
                         input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
                         attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
                         labels = batch["label"].to("cpu")
-                        optimizer.zero_grad()
                         outputs = model(input_ids)
                         loss = nn.CrossEntropyLoss()(outputs, labels)
-                        loss.backward()
-                        optimizer.step()
-                        progress_bar.update(1)
-
-                        elapsed_time = time.time() - start_time
-                        estimated_total_time = total_steps * (elapsed_time / (step + 1))
-                        estimated_remaining_time = estimated_total_time - elapsed_time
-
-                        if step % args.logging_steps == 0:
-                            train_losses.append(loss.item())
-                            print(f"Step {step}/{total_steps}, Loss: {loss.item()}, Estimated remaining time: {estimated_remaining_time:.2f} seconds")
-
-                    epoch += 1
-                    model.eval()
-                    eval_loss = 0
-                    with torch.no_grad():
-                        for batch in eval_loader:
-                            input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
-                            attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
-                            labels = batch["label"].to("cpu")
-                            outputs = model(input_ids)
-                            loss = nn.CrossEntropyLoss()(outputs, labels)
-                            eval_loss += loss.item()
-
-                    eval_loss /= len(eval_loader)
-                    eval_losses.append(eval_loss)
-                    print(f"Epoch {epoch}/{args.num_train_epochs}, Evaluation Loss: {eval_loss}")
+                        eval_loss += loss.item()
+
+                eval_loss /= len(eval_loader)
+                eval_losses.append(eval_loss)
+                print(f"Epoch {epoch}/{args.num_train_epochs}, Evaluation Loss: {eval_loss}")
 
             train(unified_model, train_loader, eval_loader, training_args)
 
@@ -257,26 +266,5 @@ def main():
             print(f"Error: {e}")
             time.sleep(2)
 
-
-
-def gradio_app():
-    with gr.Blocks() as app:
-        gr.Markdown(
-            """
-            # IA Generativa con Transformers y Diffusers
-            Explora diferentes modelos de IA para generar texto, imágenes, audio, video y más.
-            """
-        )
-    app.launch()
-
-
-
-
-
 if __name__ == "__main__":
     main()
-    gradio_app()
-
-
-
-