Update app.py

app.py

@@ -12,6 +12,7 @@ from huggingface_hub import snapshot_download, HfApi, HfFolder
 import multiprocessing
 import io
 import time
+from tqdm import tqdm
 
 hf_token = os.getenv("HF_TOKEN")
 redis_host = os.getenv("REDIS_HOST")
@@ -61,18 +62,22 @@ def get_model_or_download(model_id, redis_key, loader_func):
     if model:
         return model
     try:
-        model = loader_func(model_id, torch_dtype=torch.float16)
+        with tqdm(total=1, desc=f"Downloading {model_id}") as pbar:
+            model = loader_func(model_id, torch_dtype=torch.float16)
+            pbar.update(1)
         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):
-    redis_key = f"generated_image_{prompt}"
+    redis_key = f"generated_image:{prompt}"
     image = load_object_from_redis(redis_key)
     if not image:
         try:
-            image = text_to_image_pipeline(prompt).images[0]
+            with tqdm(total=1, desc="Generating image") as pbar:
+                image = text_to_image_pipeline(prompt).images[0]
+                pbar.update(1)
             save_object_to_redis(redis_key, image)
         except Exception as e:
             print(f"Failed to generate image: {e}")
@@ -80,11 +85,13 @@ def generate_image(prompt):
     return image
 
 def edit_image_with_prompt(image, prompt, strength=0.75):
-    redis_key = f"edited_image_{prompt}_{strength}"
+    redis_key = f"edited_image:{prompt}:{strength}"
     edited_image = load_object_from_redis(redis_key)
     if not edited_image:
         try:
-            edited_image = img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
+            with tqdm(total=1, desc="Editing image") as pbar:
+                edited_image = img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
+                pbar.update(1)
             save_object_to_redis(redis_key, edited_image)
         except Exception as e:
             print(f"Failed to edit image: {e}")
@@ -92,11 +99,13 @@ def edit_image_with_prompt(image, prompt, strength=0.75):
     return edited_image
 
 def generate_song(prompt, duration=10):
-    redis_key = f"generated_song_{prompt}_{duration}"
+    redis_key = f"generated_song:{prompt}:{duration}"
     song = load_object_from_redis(redis_key)
     if not song:
         try:
-            song = music_gen.generate(prompt, duration=duration)
+            with tqdm(total=1, desc="Generating song") as pbar:
+                song = music_gen.generate(prompt, duration=duration)
+                pbar.update(1)
             save_object_to_redis(redis_key, song)
         except Exception as e:
             print(f"Failed to generate song: {e}")
@@ -104,11 +113,13 @@ def generate_song(prompt, duration=10):
     return song
 
 def generate_text(prompt):
-    redis_key = f"generated_text_{prompt}"
+    redis_key = f"generated_text:{prompt}"
     text = load_object_from_redis(redis_key)
     if not text:
         try:
-            text = text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
+            with tqdm(total=1, desc="Generating text") as pbar:
+                text = text_gen_pipeline([{"role": "user", "content": prompt}], max_new_tokens=256)[0]["generated_text"].strip()
+                pbar.update(1)
             save_object_to_redis(redis_key, text)
         except Exception as e:
             print(f"Failed to generate text: {e}")
@@ -116,17 +127,19 @@ def generate_text(prompt):
     return text
 
 def generate_flux_image(prompt):
-    redis_key = f"generated_flux_image_{prompt}"
+    redis_key = f"generated_flux_image:{prompt}"
     flux_image = load_object_from_redis(redis_key)
     if not flux_image:
         try:
-            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]
+            with tqdm(total=1, desc="Generating FLUX image") as pbar:
+                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]
+                pbar.update(1)
             save_object_to_redis(redis_key, flux_image)
         except Exception as e:
             print(f"Failed to generate flux image: {e}")
@@ -134,13 +147,15 @@ def generate_flux_image(prompt):
     return flux_image
 
 def generate_code(prompt):
-    redis_key = f"generated_code_{prompt}"
+    redis_key = f"generated_code:{prompt}"
     code = load_object_from_redis(redis_key)
     if not code:
         try:
-            inputs = starcoder_tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-            outputs = starcoder_model.generate(inputs)
-            code = starcoder_tokenizer.decode(outputs[0])
+            with tqdm(total=1, desc="Generating code") as pbar:
+                inputs = starcoder_tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+                outputs = starcoder_model.generate(inputs)
+                code = starcoder_tokenizer.decode(outputs[0])
+                pbar.update(1)
             save_object_to_redis(redis_key, code)
         except Exception as e:
             print(f"Failed to generate code: {e}")
@@ -148,14 +163,16 @@ def generate_code(prompt):
     return code
 
 def generate_video(prompt):
-    redis_key = f"generated_video_{prompt}"
+    redis_key = f"generated_video:{prompt}"
     video = load_object_from_redis(redis_key)
     if not video:
         try:
-            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)
+            with tqdm(total=1, desc="Generating video") as pbar:
+                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)
+                pbar.update(1)
             save_object_to_redis(redis_key, video)
         except Exception as e:
             print(f"Failed to generate video: {e}")
@@ -171,7 +188,9 @@ def test_model_meta_llama():
                 {"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()
+            with tqdm(total=1, desc="Testing Meta-Llama") as pbar:
+                response = meta_llama_pipeline(messages, max_new_tokens=256)[0]["generated_text"].strip()
+                pbar.update(1)
             save_object_to_redis(redis_key, response)
         except Exception as e:
             print(f"Failed to test Meta-Llama: {e}")
@@ -188,7 +207,9 @@ def train_model(model, dataset, epochs, batch_size, learning_rate):
     )
     trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
     try:
-        trainer.train()
+        with tqdm(total=epochs, desc="Training model") as pbar:
+            trainer.train()
+            pbar.update(epochs)
         save_object_to_redis("trained_model", model)
         save_object_to_redis("training_results", output_dir.getvalue())
     except Exception as e: