Update app.py

app.py

@@ -4,12 +4,12 @@ import redis
 import torch
 import scipy
 from transformers import (
-    pipeline, AutoTokenizer, AutoModelForCausalLM, AutoProcessor, 
-    MusicgenForConditionalGeneration, WhisperProcessor, WhisperForConditionalGeneration, 
+    pipeline, AutoTokenizer, AutoModelForCausalLM, AutoProcessor,
+    MusicgenForConditionalGeneration, WhisperProcessor, WhisperForConditionalGeneration,
     MarianMTModel, MarianTokenizer, BartTokenizer, BartForConditionalGeneration
 )
 from diffusers import (
-    FluxPipeline, StableDiffusionPipeline, DPMSolverMultistepScheduler, 
+    FluxPipeline, StableDiffusionPipeline, DPMSolverMultistepScheduler,
     StableDiffusionImg2ImgPipeline, DiffusionPipeline
 )
 from diffusers.utils import export_to_video
@@ -22,16 +22,18 @@ import multiprocessing
 load_dotenv()
 
 redis_client = redis.Redis(
-    host=os.getenv('REDIS_HOST'), 
-    port=os.getenv('REDIS_PORT'), 
+    host=os.getenv('REDIS_HOST'),
+    port=os.getenv('REDIS_PORT'),
     password=os.getenv("REDIS_PASSWORD")
 )
 
 huggingface_token = os.getenv('HF_TOKEN')
 
+
 def generate_unique_id():
     return str(uuid.uuid4())
 
+
 def store_special_tokens(tokenizer, model_name):
     special_tokens = {
         'pad_token': tokenizer.pad_token,
@@ -45,6 +47,7 @@ def store_special_tokens(tokenizer, model_name):
     }
     redis_client.hmset(f"tokenizer_special_tokens:{model_name}", special_tokens)
 
+
 def load_special_tokens(tokenizer, model_name):
     special_tokens = redis_client.hgetall(f"tokenizer_special_tokens:{model_name}")
     if special_tokens:
@@ -57,6 +60,7 @@ def load_special_tokens(tokenizer, model_name):
         tokenizer.bos_token = special_tokens.get('bos_token', '').decode("utf-8")
         tokenizer.bos_token_id = int(special_tokens.get('bos_token_id', -1))
 
+
 def train_and_store_transformers_model(model_name, data):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -69,6 +73,7 @@ def train_and_store_transformers_model(model_name, data):
     tokenizer_data = tokenizer.save_pretrained("transformers_tokenizer")
     redis_client.set(f"transformers_tokenizer:{model_name}", tokenizer_data)
 
+
 def generate_transformers_response_from_redis(model_name, prompt):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"transformers_model:{model_name}:state_dict")
@@ -85,6 +90,7 @@ def generate_transformers_response_from_redis(model_name, prompt):
     redis_client.set(f"transformers_response:{unique_id}", response)
     return response
 
+
 def train_and_store_diffusers_model(model_name, data):
     pipe = FluxPipeline.from_pretrained(model_name, torch_dtype=torch.bfloat16)
     pipe.enable_model_cpu_offload()
@@ -94,6 +100,7 @@ def train_and_store_diffusers_model(model_name, data):
         model_data = f.read()
     redis_client.set(f"diffusers_model:{model_name}", model_data)
 
+
 def generate_diffusers_image_from_redis(model_name, prompt):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"diffusers_model:{model_name}")
@@ -101,12 +108,14 @@ def generate_diffusers_image_from_redis(model_name, prompt):
         f.write(model_data)
     pipe = FluxPipeline.from_pretrained("diffusers_model", torch_dtype=torch.bfloat16)
     pipe.enable_model_cpu_offload()
-    image = pipe(prompt, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, generator=torch.Generator("cpu").manual_seed(0)).images[0]
+    image = pipe(prompt, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256,
+                  generator=torch.Generator("cpu").manual_seed(0)).images[0]
     image_path = f"images/diffusers_{unique_id}.png"
     image.save(image_path)
     redis_client.set(f"diffusers_image:{unique_id}", image_path)
     return image
 
+
 def train_and_store_musicgen_model(model_name, data):
     processor = AutoProcessor.from_pretrained(model_name)
     model = MusicgenForConditionalGeneration.from_pretrained(model_name)
@@ -118,6 +127,7 @@ def train_and_store_musicgen_model(model_name, data):
     processor_data = processor.save_pretrained("musicgen_processor")
     redis_client.set(f"musicgen_processor:{model_name}", processor_data)
 
+
 def generate_musicgen_audio_from_redis(model_name, text_prompts):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"musicgen_model:{model_name}:state_dict")
@@ -134,6 +144,7 @@ def generate_musicgen_audio_from_redis(model_name, text_prompts):
     redis_client.set(f"musicgen_audio:{unique_id}", audio_path)
     return audio_path
 
+
 def train_and_store_stable_diffusion_model(model_name, data):
     pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
@@ -144,6 +155,7 @@ def train_and_store_stable_diffusion_model(model_name, data):
         model_data = f.read()
     redis_client.set(f"stable_diffusion_model:{model_name}", model_data)
 
+
 def generate_stable_diffusion_image_from_redis(model_name, prompt):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"stable_diffusion_model:{model_name}")
@@ -158,6 +170,7 @@ def generate_stable_diffusion_image_from_redis(model_name, prompt):
     redis_client.set(f"stable_diffusion_image:{unique_id}", image_path)
     return image
 
+
 def train_and_store_img2img_model(model_name, data):
     pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
     pipe = pipe.to("cuda")
@@ -167,6 +180,7 @@ def train_and_store_img2img_model(model_name, data):
         model_data = f.read()
     redis_client.set(f"img2img_model:{model_name}", model_data)
 
+
 def generate_img2img_from_redis(model_name, init_image, prompt, strength=0.75):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"img2img_model:{model_name}")
@@ -181,6 +195,7 @@ def generate_img2img_from_redis(model_name, init_image, prompt, strength=0.75):
     redis_client.set(f"img2img_image:{unique_id}", image_path)
     return image
 
+
 def train_and_store_marianmt_model(model_name, data):
     tokenizer = MarianTokenizer.from_pretrained(model_name)
     model = MarianMTModel.from_pretrained(model_name)
@@ -192,6 +207,7 @@ def train_and_store_marianmt_model(model_name, data):
     tokenizer_data = tokenizer.save_pretrained("marianmt_tokenizer")
     redis_client.set(f"marianmt_tokenizer:{model_name}", tokenizer_data)
 
+
 def translate_text_from_redis(model_name, text, src_lang, tgt_lang):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"marianmt_model:{model_name}:state_dict")
@@ -207,6 +223,7 @@ def translate_text_from_redis(model_name, text, src_lang, tgt_lang):
     redis_client.set(f"marianmt_translation:{unique_id}", translation)
     return translation
 
+
 def train_and_store_bart_model(model_name, data):
     tokenizer = BartTokenizer.from_pretrained(model_name)
     model = BartForConditionalGeneration.from_pretrained(model_name)
@@ -218,6 +235,7 @@ def train_and_store_bart_model(model_name, data):
     tokenizer_data = tokenizer.save_pretrained("bart_tokenizer")
     redis_client.set(f"bart_tokenizer:{model_name}", tokenizer_data)
 
+
 def summarize_text_from_redis(model_name, text):
     unique_id = generate_unique_id()
     model_data = redis_client.get(f"bart_model:{model_name}:state_dict")
@@ -234,6 +252,7 @@ def summarize_text_from_redis(model_name, text):
     redis_client.set(f"bart_summary:{unique_id}", summary)
     return summary
 
+
 def auto_train_and_store(model_name, task, data):
     if task == "text-generation":
         train_and_store_transformers_model(model_name, data)
@@ -250,6 +269,7 @@ def auto_train_and_store(model_name, task, data):
     elif task == "summarization":
         train_and_store_bart_model(model_name, data)
 
+
 def transcribe_audio_from_redis(audio_file):
     audio_file_path = "audio_file.wav"
     with open(audio_file_path, "wb") as f:
@@ -263,6 +283,7 @@ def transcribe_audio_from_redis(audio_file):
     transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
     return transcription[0]
 
+
 def generate_image_from_redis(model_name, prompt, model_type):
     if model_type == "diffusers":
         image = generate_diffusers_image_from_redis(model_name, prompt)
@@ -272,8 +293,10 @@ def generate_image_from_redis(model_name, prompt, model_type):
         image = generate_img2img_from_redis(model_name, "init_image.png", prompt)
     return image
 
+
 def generate_video_from_redis(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,
+                                             variant="fp16")
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
     pipe.enable_model_cpu_offload()
     video_frames = pipe(prompt, num_inference_steps=25).frames
@@ -282,6 +305,7 @@ def generate_video_from_redis(prompt):
     redis_client.set(f"video_{unique_id}", video_path)
     return video_path
 
+
 def generate_random_response(prompts, generator):
     responses = []
     for prompt in prompts:
@@ -289,16 +313,19 @@ def generate_random_response(prompts, generator):
         responses.append(response)
     return responses
 
+
 def process_parallel(tasks):
     with multiprocessing.Pool() as pool:
         results = pool.map(lambda task: task(), tasks)
     return results
 
+
 def generate_response_from_prompt(prompt, model_name="google/flan-t5-xl"):
     generator = pipeline('text-generation', model=model_name, tokenizer=model_name)
     responses = generate_random_response([prompt], generator)
     return responses[0]
 
+
 def generate_image_from_prompt(prompt, image_type, model_name="CompVis/stable-diffusion-v1-4"):
     if image_type == "diffusers":
         image = generate_diffusers_image_from_redis(model_name, prompt)
@@ -308,65 +335,98 @@ def generate_image_from_prompt(prompt, image_type, model_name="CompVis/stable-di
         image = generate_img2img_from_redis(model_name, "init_image.png", prompt)
     return image
 
+
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("## Generación de Texto con Transformers")
-        with gr.Row():
-            prompt_text = gr.Textbox(label="Texto de Entrada")
-            text_output = gr.Textbox(label="Respuesta")
-            text_button = gr.Button("Generar Texto")
-        text_button.click(generate_response_from_prompt, inputs=prompt_text, outputs=text_output)
-
-        gr.Markdown("## Generación de Imágenes con Diffusers, Stable Diffusion e Img2Img")
-        with gr.Row():
-            prompt_image = gr.Textbox(label="Prompt de Imagen")
-            image_type = gr.Dropdown(["diffusers", "stable-diffusion", "img2img"], label="Tipo de Imagen")
-            model_name_image = gr.Textbox(label="Nombre del Modelo", value="CompVis/stable-diffusion-v1-4")
-            image_output = gr.Image(label="Imagen Generada")
-            image_button = gr.Button("Generar Imagen")
-        image_button.click(generate_image_from_prompt, inputs=[prompt_image, image_type, model_name_image], outputs=image_output)
-
-        gr.Markdown("## Generación de Video")
-        with gr.Row():
-            prompt_video = gr.Textbox(label="Prompt de Video")
-            video_output = gr.Video(label="Video Generado")
-            video_button = gr.Button("Generar Video")
-        video_button.click(generate_video_from_redis, inputs=prompt_video, outputs=video_output)
-
-        gr.Markdown("## Generación de Audio con MusicGen")
-        with gr.Row():
-            model_name_audio = gr.Textbox(label="Nombre del Modelo", value="facebook/musicgen-small")
-            text_prompts_audio = gr.Textbox(label="Prompts de Audio")
-            audio_output = gr.Audio(label="Audio Generado")
-            audio_button = gr.Button("Generar Audio")
-        audio_button.click(generate_musicgen_audio_from_redis, inputs=[model_name_audio, text_prompts_audio], outputs=audio_output)
-
-        gr.Markdown("## Transcripción de Audio con Whisper")
-        with gr.Row():
-            audio_file = gr.Audio(type="filepath", label="Archivo de Audio")
-            transcription_output = gr.Textbox(label="Transcripción")
-            audio_button = gr.Button("Transcribir Audio")
-        audio_button.click(transcribe_audio_from_redis, inputs=audio_file, outputs=transcription_output)
-
-        gr.Markdown("## Traducción de Texto")
-        with gr.Row():
-            model_name_translate = gr.Textbox(label="Nombre del Modelo", value="Helsinki-NLP/opus-mt-en-es")
-            text_input = gr.Textbox(label="Texto a Traducir")
-            translation_output = gr.Textbox(label="Traducción")
-            src_lang_input = gr.Textbox(label="Idioma de Origen", value="en")
-            tgt_lang_input = gr.Textbox(label="Idioma de Destino", value="es")
-            translate_button = gr.Button("Traducir Texto")
-        translate_button.click(translate_text_from_redis, inputs=[model_name_translate, text_input, src_lang_input, tgt_lang_input], outputs=translation_output)
-
-        gr.Markdown("## Resumen de Texto")
-        with gr.Row():
-            model_name_summarize = gr.Textbox(label="Nombre del Modelo", value="facebook/bart-large-cnn")
-            text_to_summarize = gr.Textbox(label="Texto para Resumir")
-            summary_output = gr.Textbox(label="Resumen")
-            summarize_button = gr.Button("Generar Resumen")
-        summarize_button.click(summarize_text_from_redis, inputs=[model_name_summarize, text_to_summarize], outputs=summary_output)
+        gr.Markdown(
+            """
+            # IA Generativa con Transformers y Diffusers
+            Explora diferentes modelos de IA para generar texto, imágenes, audio, video y más.
+            """
+        )
+
+        with gr.Tab("Texto"):
+            with gr.Row():
+                with gr.Column():
+                    prompt_text = gr.Textbox(label="Texto de Entrada", placeholder="Ingresa tu prompt de texto aquí...")
+                    text_button = gr.Button("Generar Texto", variant="primary")
+                with gr.Column():
+                    text_output = gr.Textbox(label="Respuesta")
+            text_button.click(generate_response_from_prompt, inputs=prompt_text, outputs=text_output)
+
+        with gr.Tab("Imagen"):
+            with gr.Row():
+                with gr.Column():
+                    prompt_image = gr.Textbox(label="Prompt de Imagen",
+                                             placeholder="Ingresa tu prompt de imagen aquí...")
+                    image_type = gr.Dropdown(["diffusers", "stable-diffusion", "img2img"], label="Tipo de Modelo",
+                                            value="stable-diffusion")
+                    model_name_image = gr.Textbox(label="Nombre del Modelo",
+                                                  value="CompVis/stable-diffusion-v1-4")
+                    image_button = gr.Button("Generar Imagen", variant="primary")
+                with gr.Column():
+                    image_output = gr.Image(label="Imagen Generada")
+            image_button.click(generate_image_from_prompt, inputs=[prompt_image, image_type, model_name_image],
+                               outputs=image_output)
+
+        with gr.Tab("Video"):
+            with gr.Row():
+                with gr.Column():
+                    prompt_video = gr.Textbox(label="Prompt de Video", placeholder="Ingresa tu prompt de video aquí...")
+                    video_button = gr.Button("Generar Video", variant="primary")
+                with gr.Column():
+                    video_output = gr.Video(label="Video Generado")
+            video_button.click(generate_video_from_redis, inputs=prompt_video, outputs=video_output)
+
+        with gr.Tab("Audio"):
+            with gr.Row():
+                with gr.Column():
+                    model_name_audio = gr.Textbox(label="Nombre del Modelo", value="facebook/musicgen-small")
+                    text_prompts_audio = gr.Textbox(label="Prompts de Audio",
+                                                   placeholder="Ingresa tus prompts de audio aquí...")
+                    audio_button = gr.Button("Generar Audio", variant="primary")
+                with gr.Column():
+                    audio_output = gr.Audio(label="Audio Generado")
+            audio_button.click(generate_musicgen_audio_from_redis, inputs=[model_name_audio, text_prompts_audio],
+                               outputs=audio_output)
+
+        with gr.Tab("Transcripción"):
+            with gr.Row():
+                with gr.Column():
+                    audio_file = gr.Audio(type="filepath", label="Archivo de Audio")
+                    audio_button = gr.Button("Transcribir Audio", variant="primary")
+                with gr.Column():
+                    transcription_output = gr.Textbox(label="Transcripción")
+            audio_button.click(transcribe_audio_from_redis, inputs=audio_file, outputs=transcription_output)
+
+        with gr.Tab("Traducción"):
+            with gr.Row():
+                with gr.Column():
+                    model_name_translate = gr.Textbox(label="Nombre del Modelo", value="Helsinki-NLP/opus-mt-en-es")
+                    text_input = gr.Textbox(label="Texto a Traducir", placeholder="Ingresa el texto a traducir...")
+                    src_lang_input = gr.Textbox(label="Idioma de Origen", value="en")
+                    tgt_lang_input = gr.Textbox(label="Idioma de Destino", value="es")
+                    translate_button = gr.Button("Traducir Texto", variant="primary")
+                with gr.Column():
+                    translation_output = gr.Textbox(label="Traducción")
+            translate_button.click(translate_text_from_redis,
+                                   inputs=[model_name_translate, text_input, src_lang_input, tgt_lang_input],
+                                   outputs=translation_output)
+
+        with gr.Tab("Resumen"):
+            with gr.Row():
+                with gr.Column():
+                    model_name_summarize = gr.Textbox(label="Nombre del Modelo", value="facebook/bart-large-cnn")
+                    text_to_summarize = gr.Textbox(label="Texto para Resumir",
+                                                  placeholder="Ingresa el texto a resumir...")
+                    summarize_button = gr.Button("Generar Resumen", variant="primary")
+                with gr.Column():
+                    summary_output = gr.Textbox(label="Resumen")
+            summarize_button.click(summarize_text_from_redis, inputs=[model_name_summarize, text_to_summarize],
+                                   outputs=summary_output)
 
     app.launch()
 
+
 if __name__ == "__main__":
     gradio_app()