Update app.py

app.py

@@ -5,17 +5,18 @@ 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, TrainingArguments, Trainer
+from transformers import pipeline as transformers_pipeline, AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from audiocraft.models import MusicGen
 import gradio as gr
-from huggingface_hub import HfFolder
+from huggingface_hub import snapshot_download, HfApi, HfFolder
 import multiprocessing
 import io
 import time
 
+# Obtener las variables de entorno
 hf_token = os.getenv("HF_TOKEN")
 redis_host = os.getenv("REDIS_HOST")
-redis_port = int(os.getenv("REDIS_PORT", 6379))
+redis_port = int(os.getenv("REDIS_PORT", 6379))  # Valor predeterminado si no se proporciona
 redis_password = os.getenv("REDIS_PASSWORD")
 
 HfFolder.save_token(hf_token)
@@ -24,15 +25,18 @@ def connect_to_redis():
     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"Connection to Redis failed: {e}. Retrying in 1 second...")
             time.sleep(1)
 
 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
 
@@ -43,6 +47,7 @@ def load_object_from_redis(key):
         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):
@@ -50,17 +55,21 @@ def save_object_to_redis(key, obj):
     redis_client = reconnect_if_needed(redis_client)
     try:
         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 model:
+        print(f"Model loaded from Redis: {redis_key}")
         return model
     try:
         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}")
     except Exception as e:
+        print(f"Failed to load or save model: {e}")
         return None
 
 def generate_image(prompt):
@@ -71,6 +80,7 @@ def generate_image(prompt):
             image = text_to_image_pipeline(prompt).images[0]
             save_object_to_redis(redis_key, image)
         except Exception as e:
+            print(f"Failed to generate image: {e}")
             return None
     return image
 
@@ -82,6 +92,7 @@ def edit_image_with_prompt(image, prompt, strength=0.75):
             edited_image = img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
             save_object_to_redis(redis_key, edited_image)
         except Exception as e:
+            print(f"Failed to edit image: {e}")
             return None
     return edited_image
 
@@ -93,6 +104,7 @@ def generate_song(prompt, duration=10):
             song = music_gen.generate(prompt, duration=duration)
             save_object_to_redis(redis_key, song)
         except Exception as e:
+            print(f"Failed to generate song: {e}")
             return None
     return song
 
@@ -101,10 +113,10 @@ def generate_text(prompt):
     text = load_object_from_redis(redis_key)
     if not text:
         try:
-            # Reemplazar "bigcode/starcoder" con otro modelo de generación de texto
             text = text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
             save_object_to_redis(redis_key, text)
         except Exception as e:
+            print(f"Failed to generate text: {e}")
             return None
     return text
 
@@ -122,6 +134,7 @@ def generate_flux_image(prompt):
             ).images[0]
             save_object_to_redis(redis_key, flux_image)
         except Exception as e:
+            print(f"Failed to generate flux image: {e}")
             return None
     return flux_image
 
@@ -135,6 +148,7 @@ def generate_code(prompt):
             code = starcoder_tokenizer.decode(outputs[0])
             save_object_to_redis(redis_key, code)
         except Exception as e:
+            print(f"Failed to generate code: {e}")
             return None
     return code
 
@@ -149,6 +163,7 @@ def generate_video(prompt):
             video = export_to_video(pipe(prompt, num_inference_steps=25).frames)
             save_object_to_redis(redis_key, video)
         except Exception as e:
+            print(f"Failed to generate video: {e}")
             return None
     return video
 
@@ -164,6 +179,7 @@ def test_model_meta_llama():
             response = meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"].strip()
             save_object_to_redis(redis_key, response)
         except Exception as e:
+            print(f"Failed to test Meta-Llama: {e}")
             return None
     return response
 
@@ -205,13 +221,15 @@ for _ in range(num_processes):
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Cargar modelos
 text_to_image_pipeline = get_model_or_download("stabilityai/stable-diffusion-2", "text_to_image_model", StableDiffusionPipeline.from_pretrained)
 img2img_pipeline = get_model_or_download("CompVis/stable-diffusion-v1-4", "img2img_model", StableDiffusionImg2ImgPipeline.from_pretrained)
 flux_pipeline = get_model_or_download("black-forest-labs/FLUX.1-schnell", "flux_model", FluxPipeline.from_pretrained)
-text_gen_pipeline = transformers_pipeline("text-generation", model="google/flan-t5-xl", tokenizer="google/flan-t5-xl", device=device)
+text_gen_pipeline = transformers_pipeline("text-generation", model="google/gemma-2-9b", tokenizer="google/gemma-2-9b", device=0)
 music_gen = load_object_from_redis("music_gen") or MusicGen.from_pretrained('melody')
 meta_llama_pipeline = get_model_or_download("meta-llama/Meta-Llama-3.1-8B-Instruct", "meta_llama_model", transformers_pipeline)
 
+# Definir interfaces de usuario
 gen_image_tab = gr.Interface(generate_image, gr.inputs.Textbox(label="Prompt:"), gr.outputs.Image(type="pil"), title="Generate Image")
 edit_image_tab = gr.Interface(edit_image_with_prompt, [gr.inputs.Image(type="pil", label="Image:"), gr.inputs.Textbox(label="Prompt:"), gr.inputs.Slider(0.1, 1.0, 0.75, step=0.05, label="Strength:")], gr.outputs.Image(type="pil"), title="Edit Image")
 generate_song_tab = gr.Interface(generate_song, [gr.inputs.Textbox(label="Prompt:"), gr.inputs.Slider(5, 60, 10, step=1, label="Duration (s):")], gr.outputs.Audio(type="numpy"), title="Generate Songs")
@@ -229,4 +247,4 @@ app.launch(share=True)
 for _ in range(num_processes):
     task_queue.put(None)
 for p in processes:
-    p.join()
+    p.join()