Update app.py

app.py

@@ -1,20 +1,18 @@
+!pip install redis diffusers transformers accelerate torch gradio audiocraft huggingface_hub
+
 import redis
 import pickle
 import torch
 from PIL import Image
 from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, FluxPipeline, DiffusionPipeline, DPMSolverMultistepScheduler
 from diffusers.utils import export_to_video
-from transformers import pipeline as transformers_pipeline, AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+from transformers import pipeline as transformers_pipeline, AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from audiocraft.models import MusicGen
 import gradio as gr
 from huggingface_hub import snapshot_download, HfApi, HfFolder
 import multiprocessing
 import io
-from dotenv import load_dotenv
-import os
-
-# Cargar las variables del archivo .env
-load_dotenv()
+import time
 
 # Obtener las variables de entorno
 hf_token = os.getenv("HF_TOKEN")
@@ -22,26 +20,44 @@ redis_host = os.getenv("REDIS_HOST")
 redis_port = os.getenv("REDIS_PORT")
 redis_password = os.getenv("REDIS_PASSWORD")
 
-# Usar las variables de huggingface
 HfFolder.save_token(hf_token)
 
-# Usar las variables de redis
 def connect_to_redis():
-    return  redis.Redis(host=redis_host, port=redis_port, password=redis_password)
+    max_retries = 5
+    retry_delay = 1
+    for attempt in range(max_retries):
+        try:
+            redis_client = redis.Redis(host=redis_host, port=redis_port, password=redis_password)
+            redis_client.ping()
+            return redis_client
+        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError, BrokenPipeError) as e:
+            print(f"Attempt {attempt + 1}: Connection to Redis failed: {e}. Retrying in {retry_delay} seconds...")
+            time.sleep(retry_delay)
+    raise ConnectionError("Failed to connect to Redis after multiple retries.")
+
+def reconnect_if_needed(redis_client):
+    try:
+        redis_client.ping()
+    except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError, BrokenPipeError):
+        print("Reconnecting to Redis...")
+        return connect_to_redis()
+    return redis_client
 
 def load_object_from_redis(key):
-    with connect_to_redis() as redis_client:
-        obj_data = redis_client.get(key)
+    redis_client = connect_to_redis()
+    redis_client = reconnect_if_needed(redis_client)
+    obj_data = redis_client.get(key)
     return pickle.loads(obj_data) if obj_data else None
 
 def save_object_to_redis(key, obj):
-    with connect_to_redis() as redis_client:
-        redis_client.set(key, pickle.dumps(obj))
+    redis_client = connect_to_redis()
+    redis_client = reconnect_if_needed(redis_client)
+    redis_client.set(key, pickle.dumps(obj))
 
 def get_model_or_download(model_id, redis_key, loader_func):
     model = load_object_from_redis(redis_key)
     if not model:
-        model = loader_func(model_id, use_auth_token=hf_token, torch_dtype=torch.float16)
+        model = loader_func(model_id, torch_dtype=torch.float16)
         save_object_to_redis(redis_key, model)
     return model
 
@@ -55,7 +71,7 @@ def generate_song(prompt, duration=10):
     return music_gen.generate(prompt, duration=duration)
 
 def generate_text(prompt):
-    return text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"][-1]["content"].strip()
+    return text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
 
 def generate_flux_image(prompt):
     return flux_pipeline(
@@ -72,7 +88,7 @@ def generate_code(prompt):
     return starcoder_tokenizer.decode(outputs[0])
 
 def generate_video(prompt):
-    pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16")
+    pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16)
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     pipe.enable_model_cpu_offload()
     return export_to_video(pipe(prompt, num_inference_steps=25).frames)
@@ -82,7 +98,7 @@ def test_model_meta_llama():
         {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
         {"role": "user", "content": "Who are you?"}
     ]
-    return meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"][-1]
+    return meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"].strip()
 
 def train_model(model, dataset, epochs, batch_size, learning_rate):
     output_dir = io.BytesIO()
@@ -156,4 +172,4 @@ app.launch(share=True)
 for _ in range(num_processes):
     task_queue.put(None)
 for p in processes:
-    p.join()
+    p.join()