Update app.py

app.py

@@ -6,17 +6,16 @@ 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, Trainer
-from audiocraft.models import MusicGen
+from audiocraft.models import musicgen
 import gradio as gr
 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))  # Valor predeterminado si no se proporciona
+redis_port = int(os.getenv("REDIS_PORT", 6379))
 redis_password = os.getenv("REDIS_PASSWORD")
 
 HfFolder.save_token(hf_token)
@@ -25,8 +24,7 @@ def connect_to_redis():
     while True:
         try:
             redis_client = redis.Redis(host=redis_host, port=redis_port, password=redis_password)
-            redis_client.ping()  # Verifica si la conexión está activa
-            print("Connected to Redis successfully.")
+            redis_client.ping()
             return redis_client
         except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError, BrokenPipeError) as e:
             print(f"Connection to Redis failed: {e}. Retrying in 1 second...")
@@ -55,19 +53,16 @@ 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
@@ -221,25 +216,26 @@ 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/gemma-2-9b", tokenizer="google/gemma-2-9b", device=0)
-music_gen = load_object_from_redis("music_gen") or MusicGen.from_pretrained('melody')
+music_gen = load_object_from_redis("music_gen") or musicgen.MusicGen.get_pretrained('melody')
 meta_llama_pipeline = get_model_or_download("meta-llama/Meta-Llama-3.1-8B-Instruct", "meta_llama_model", transformers_pipeline)
+starcoder_model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder")
+starcoder_tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoder")
 
-# 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")
 generate_text_tab = gr.Interface(generate_text, gr.inputs.Textbox(label="Prompt:"), gr.outputs.Textbox(label="Generated Text:"), title="Generate Text")
 generate_flux_image_tab = gr.Interface(generate_flux_image, gr.inputs.Textbox(label="Prompt:"), gr.outputs.Image(type="pil"), title="Generate FLUX Images")
+generate_code_tab = gr.Interface(generate_code, gr.inputs.Textbox(label="Prompt:"), gr.outputs.Textbox(label="Generated Code:"), title="Generate Code")
 model_meta_llama_test_tab = gr.Interface(test_model_meta_llama, gr.inputs.Textbox(label="Test Input:"), gr.outputs.Textbox(label="Model Output:"), title="Test Meta-Llama")
 
 app = gr.TabbedInterface(
-    [gen_image_tab, edit_image_tab, generate_song_tab, generate_text_tab, generate_flux_image_tab, model_meta_llama_test_tab],
-    ["Generate Image", "Edit Image", "Generate Song", "Generate Text", "Generate FLUX Image", "Test Meta-Llama"]
+    [gen_image_tab, edit_image_tab, generate_song_tab, generate_text_tab, generate_flux_image_tab, generate_code_tab, model_meta_llama_test_tab],
+    ["Generate Image", "Edit Image", "Generate Song", "Generate Text", "Generate FLUX Image", "Generate Code", "Test Meta-Llama"]
 )
 
 app.launch(share=True)
@@ -247,4 +243,4 @@ app.launch(share=True)
 for _ in range(num_processes):
     task_queue.put(None)
 for p in processes:
-    p.join()
+    p.join()