Update app.py

app.py

@@ -1,5 +1,3 @@
-!pip install redis diffusers transformers accelerate torch gradio audiocraft huggingface_hub
-
 import redis
 import pickle
 import torch
@@ -23,17 +21,15 @@ redis_password = os.getenv("REDIS_PASSWORD")
 HfFolder.save_token(hf_token)
 
 def connect_to_redis():
-    max_retries = 5
-    retry_delay = 1
-    for attempt in range(max_retries):
+    while True:
         try:
             redis_client = redis.Redis(host=redis_host, port=redis_port, password=redis_password)
-            redis_client.ping()
+            redis_client.ping()  # Verifica si la conexión está activa
+            print("Connected to Redis successfully.")
             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.")
+            print(f"Connection to Redis failed: {e}. Retrying in 1 second...")
+            time.sleep(1)
 
 def reconnect_if_needed(redis_client):
     try:
@@ -46,59 +42,111 @@ def reconnect_if_needed(redis_client):
 def load_object_from_redis(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
+    try:
+        obj_data = redis_client.get(key)
+        return pickle.loads(obj_data) if obj_data else None
+    except (pickle.PickleError, redis.exceptions.RedisError) as e:
+        print(f"Failed to load object from Redis: {e}")
+        return None
 
 def save_object_to_redis(key, obj):
     redis_client = connect_to_redis()
     redis_client = reconnect_if_needed(redis_client)
-    redis_client.set(key, pickle.dumps(obj))
+    try:
+        if not redis_client.exists(key):  # Solo guarda si no existe
+            redis_client.set(key, pickle.dumps(obj))
+            print(f"Object saved to Redis: {key}")
+    except redis.exceptions.RedisError as e:
+        print(f"Failed to save object to Redis: {e}")
 
 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, torch_dtype=torch.float16)
-        save_object_to_redis(redis_key, model)
+    if model:
+        print(f"Model loaded from Redis: {redis_key}")
+        return model
+    model = loader_func(model_id, torch_dtype=torch.float16)
+    save_object_to_redis(redis_key, model)
+    print(f"Model downloaded and saved to Redis: {redis_key}")
     return model
 
 def generate_image(prompt):
-    return text_to_image_pipeline(prompt).images[0]
+    redis_key = f"generated_image_{prompt}"
+    image = load_object_from_redis(redis_key)
+    if not image:
+        image = text_to_image_pipeline(prompt).images[0]
+        save_object_to_redis(redis_key, image)
+    return image
 
 def edit_image_with_prompt(image, prompt, strength=0.75):
-    return img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
+    redis_key = f"edited_image_{prompt}_{strength}"
+    edited_image = load_object_from_redis(redis_key)
+    if not edited_image:
+        edited_image = img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
+        save_object_to_redis(redis_key, edited_image)
+    return edited_image
 
 def generate_song(prompt, duration=10):
-    return music_gen.generate(prompt, duration=duration)
+    redis_key = f"generated_song_{prompt}_{duration}"
+    song = load_object_from_redis(redis_key)
+    if not song:
+        song = music_gen.generate(prompt, duration=duration)
+        save_object_to_redis(redis_key, song)
+    return song
 
 def generate_text(prompt):
-    return text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
+    redis_key = f"generated_text_{prompt}"
+    text = load_object_from_redis(redis_key)
+    if not text:
+        text = text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
+        save_object_to_redis(redis_key, text)
+    return text
 
 def generate_flux_image(prompt):
-    return flux_pipeline(
-        prompt,
-        guidance_scale=0.0,
-        num_inference_steps=4,
-        max_sequence_length=256,
-        generator=torch.Generator("cpu").manual_seed(0)
-    ).images[0]
+    redis_key = f"generated_flux_image_{prompt}"
+    flux_image = load_object_from_redis(redis_key)
+    if not flux_image:
+        flux_image = flux_pipeline(
+            prompt,
+            guidance_scale=0.0,
+            num_inference_steps=4,
+            max_sequence_length=256,
+            generator=torch.Generator("cpu").manual_seed(0)
+        ).images[0]
+        save_object_to_redis(redis_key, flux_image)
+    return flux_image
 
 def generate_code(prompt):
-    inputs = starcoder_tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-    outputs = starcoder_model.generate(inputs)
-    return starcoder_tokenizer.decode(outputs[0])
+    redis_key = f"generated_code_{prompt}"
+    code = load_object_from_redis(redis_key)
+    if not code:
+        inputs = starcoder_tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+        outputs = starcoder_model.generate(inputs)
+        code = starcoder_tokenizer.decode(outputs[0])
+        save_object_to_redis(redis_key, code)
+    return code
 
 def generate_video(prompt):
-    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)
+    redis_key = f"generated_video_{prompt}"
+    video = load_object_from_redis(redis_key)
+    if not video:
+        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()
+        video = export_to_video(pipe(prompt, num_inference_steps=25).frames)
+        save_object_to_redis(redis_key, video)
+    return video
 
 def test_model_meta_llama():
-    messages = [
-        {"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"].strip()
+    redis_key = "meta_llama_test_response"
+    response = load_object_from_redis(redis_key)
+    if not response:
+        messages = [
+            {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
+            {"role": "user", "content": "Who are you?"}
+        ]
+        response = meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"].strip()
+        save_object_to_redis(redis_key, response)
+    return response
 
 def train_model(model, dataset, epochs, batch_size, learning_rate):
     output_dir = io.BytesIO()
@@ -139,7 +187,7 @@ music_gen = load_object_from_redis("music_gen") or MusicGen.get_pretrained('melo
 save_object_to_redis("music_gen", music_gen)
 text_gen_pipeline = load_object_from_redis("text_gen_pipeline") or transformers_pipeline(
     "text-generation",
-    model="google/gemma-2-2b-it",
+    model="google/gemini-2-2b-it",
     model_kwargs={"torch_dtype": torch.bfloat16},
     device=device,
     use_auth_token=hf_token,