Update app.py

app.py

@@ -5,10 +5,10 @@ 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, Trainer
+from transformers import pipeline as transformers_pipeline, TrainingArguments, Trainer
 from audiocraft.models import MusicGen
 import gradio as gr
-from huggingface_hub import snapshot_download, HfApi, HfFolder
+from huggingface_hub import HfFolder
 import multiprocessing
 import io
 import time
@@ -27,14 +27,12 @@ def connect_to_redis():
             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...")
             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
 
@@ -45,7 +43,6 @@ 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):
@@ -64,7 +61,6 @@ def get_model_or_download(model_id, redis_key, loader_func):
         model = loader_func(model_id, torch_dtype=torch.float16)
         save_object_to_redis(redis_key, model)
     except Exception as e:
-        print(f"Failed to load or save model: {e}")
         return None
 
 def generate_image(prompt):
@@ -75,7 +71,6 @@ 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
 
@@ -87,7 +82,6 @@ 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
 
@@ -99,7 +93,6 @@ 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
 
@@ -108,10 +101,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
 
@@ -129,7 +122,6 @@ 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
 
@@ -143,7 +135,6 @@ 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
 
@@ -158,7 +149,6 @@ 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
 
@@ -174,7 +164,6 @@ 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
 
@@ -219,7 +208,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 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="bigcode/starcoder", tokenizer="bigcode/starcoder", device=0)
+text_gen_pipeline = transformers_pipeline("text-generation", model="google/flan-t5-xl", tokenizer="google/flan-t5-xl", device=device)
 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)