Update main.py

main.py

@@ -8,13 +8,17 @@ from transformers import (
     AutoModelForCausalLM,
     TrainingArguments,
     Trainer,
+    AutoModelForTextToWaveform,
+    pipeline,
 )
+from diffusers import FluxPipeline
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import HTMLResponse
 import multiprocessing
 import uuid
 import torch
 from torch.utils.data import Dataset
+import numpy as np
 
 load_dotenv()
 
@@ -101,6 +105,10 @@ conversation_history = {}
 tokenizer_name = "unified_tokenizer"
 tokenizer = None
 unified_model = None
+musicgen_tokenizer = AutoTokenizer.from_pretrained("facebook/musicgen-small")
+musicgen_model = AutoModelForTextToWaveform.from_pretrained("facebook/musicgen-small")
+image_pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+image_pipeline.enable_model_cpu_offload()
 
 @app.on_event("startup")
 async def startup_event():
@@ -318,8 +326,85 @@ def train_unified_model():
             trainer.train()
             unified_model.save_pretrained(model_path)
 
+async def auto_learn():
+    global tokenizer, unified_model
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD)
+    while True:
+        training_data = redis_client.lpop("training_queue")
+        if training_data:
+            item_data = json.loads(training_data)
+            tokenizer_data = item_data["tokenizers"]
+            tokenizer_name = list(tokenizer_data.keys())[0]
+            if redis_client.exists(f"tokenizer:{tokenizer_name}"):
+                tokenizer.add_tokens(list(tokenizer_data[tokenizer_name].keys()))
+            data = item_data["data"]
+            dataset = SyntheticDataset(tokenizer, data)
+
+            model_name = "unified_model"
+            model_path = f"models/{model_name}"
+
+            training_args = TrainingArguments(
+                output_dir="./results",
+                per_device_train_batch_size=8,
+                num_train_epochs=3,
+            )
+            trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset)
+            trainer.train()
+            unified_model.save_pretrained(model_path)
+
+async def auto_learn_music():
+    global musicgen_tokenizer, musicgen_model
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD)
+    while True:
+        music_training_data = redis_client.lpop("music_training_queue")
+        if music_training_data:
+            music_training_data = json.loads(music_training_data.decode("utf-8"))
+            inputs = musicgen_tokenizer(music_training_data, return_tensors="pt", padding=True)
+            musicgen_model.train()
+            optimizer = torch.optim.Adam(musicgen_model.parameters(), lr=5e-5)
+            loss_fn = torch.nn.CrossEntropyLoss()
+
+            for epoch in range(1):
+                outputs = musicgen_model(**inputs)
+                loss = loss_fn(outputs.logits, inputs['labels'])
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+async def auto_learn_images():
+    global image_pipeline
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD)
+    while True:
+        image_training_data = redis_client.lpop("image_training_queue")
+        if image_training_data:
+            image_training_data = json.loads(image_training_data.decode("utf-8"))
+            for image_prompt in image_training_data:
+                image = image_pipeline(
+                    image_prompt,
+                    guidance_scale=0.0,
+                    num_inference_steps=4,
+                    max_sequence_length=256,
+                    generator=torch.Generator("cpu").manual_seed(0)
+                ).images[0]
+                image_tensor = torch.tensor(np.array(image)).unsqueeze(0)
+                image_pipeline.model.train()
+                optimizer = torch.optim.Adam(image_pipeline.model.parameters(), lr=1e-5)
+                loss_fn = torch.nn.MSELoss()
+                target_tensor = torch.zeros_like(image_tensor)
+                for epoch in range(1):
+                    outputs = image_pipeline.model(image_tensor)
+                    loss = loss_fn(outputs, target_tensor)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+
+
 if __name__ == "__main__":
     training_process = multiprocessing.Process(target=train_unified_model)
     training_process.start()
+    music_training_process = multiprocessing.Process(target=auto_learn_music)
+    music_training_process.start()
+    image_training_process = multiprocessing.Process(target=auto_learn_images)
+    image_training_process.start()
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)