Update app.py

app.py

@@ -5,7 +5,7 @@ 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, AutoModelForCausalLM, AutoTokenizer
 from audiocraft.models import musicgen
 import gradio as gr
 from huggingface_hub import snapshot_download, HfApi, HfFolder
@@ -84,6 +84,7 @@ def get_model_or_download(model_id, redis_key, loader_func):
         save_object_to_redis(redis_key, model)
         model_bytes = pickle.dumps(model)
         upload_to_gcs(gcs_bucket_name, redis_key, model_bytes)
+        return model
     except Exception as e:
         print(f"Failed to load or save model: {e}")
         return None
@@ -113,7 +114,7 @@ def edit_image_with_prompt(image_bytes, prompt, strength=0.75):
         try:
             image = Image.open(io.BytesIO(image_bytes))
             with tqdm(total=1, desc="Editing image") as pbar:
-                edited_image = img2img_pipeline(prompt=prompt, init_image=image.convert("RGB"), strength=strength).images[0]
+                edited_image = img2img_pipeline(prompt=prompt, image=image, strength=strength).images[0]
                 pbar.update(1)
             buffered = io.BytesIO()
             edited_image.save(buffered, format="JPEG")
@@ -131,7 +132,7 @@ def generate_song(prompt, duration=10):
     if not song_bytes:
         try:
             with tqdm(total=1, desc="Generating song") as pbar:
-                song = music_gen.generate([prompt], duration=[duration])
+                song = music_gen(prompt, duration=duration)
                 pbar.update(1)
             song_bytes = song[0].getvalue()
             save_object_to_redis(redis_key, song_bytes)
@@ -166,7 +167,7 @@ def generate_flux_image(prompt):
                     prompt,
                     guidance_scale=0.0,
                     num_inference_steps=4,
-                    max_sequence_length=256,
+                    max_length=256,
                     generator=torch.Generator("cpu").manual_seed(0)
                 ).images[0]
                 pbar.update(1)
@@ -197,24 +198,6 @@ def generate_code(prompt):
             return None
     return code
 
-def generate_video(prompt):
-    redis_key = f"generated_video:{prompt}"
-    video = load_object_from_redis(redis_key)
-    if not video:
-        try:
-            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)
-            upload_to_gcs(gcs_bucket_name, redis_key, video.encode())
-        except Exception as e:
-            print(f"Failed to generate video: {e}")
-            return None
-    return video
-
 def test_model_meta_llama():
     redis_key = "meta_llama_test_response"
     response = load_object_from_redis(redis_key)
@@ -234,55 +217,15 @@ def test_model_meta_llama():
             return None
     return response
 
-def train_model(model, dataset, epochs, batch_size, learning_rate):
-    output_dir = io.BytesIO()
-    training_args = TrainingArguments(
-        output_dir=output_dir,
-        num_train_epochs=epochs,
-        per_device_train_batch_size=batch_size,
-        learning_rate=learning_rate,
-    )
-    trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
-    try:
-        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())
-        upload_to_gcs(gcs_bucket_name, "trained_model", pickle.dumps(model))
-        upload_to_gcs(gcs_bucket_name, "training_results", output_dir.getvalue())
-    except Exception as e:
-        print(f"Failed to train model: {e}")
-
-def run_task(task_queue):
-    while True:
-        task = task_queue.get()
-        if task is None:
-            break
-        func, args, kwargs = task
-        try:
-            func(*args, **kwargs)
-        except Exception as e:
-            print(f"Failed to run task: {e}")
-
-task_queue = multiprocessing.Queue()
-num_processes = multiprocessing.cpu_count()
-
-processes = []
-for _ in range(num_processes):
-    p = multiprocessing.Process(target=run_task, args=(task_queue,))
-    p.start()
-    processes.append(p)
-
 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="google/gemma-2-9b", tokenizer="google/gemma-2-9b")
-music_gen = load_object_from_redis("music_gen") or musicgen.MusicGen.get_pretrained('melody')
+music_gen = load_object_from_redis("music_gen") or musicgen.MusicGen.get_pretrained('melody').to(device)
 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_model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder").to(device)
 starcoder_tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoder")
 
 gen_image_tab = gr.Interface(fn=generate_image, inputs=gr.Textbox(label="Prompt:"), outputs=gr.Image(type="pil"), title="Generate Image")
@@ -298,9 +241,4 @@ app = gr.TabbedInterface(
     ["Generate Image", "Edit Image", "Generate Song", "Generate Text", "Generate FLUX Image", "Generate Code", "Test Meta-Llama"]
 )
 
-app.launch(share=True)
-
-for _ in range(num_processes):
-    task_queue.put(None)
-for p in processes:
-    p.join()
+app.launch(share=True)