Update app.py

app.py

@@ -72,31 +72,38 @@ def get_model_or_download(model_id, redis_key, loader_func):
 
 def generate_image(prompt):
     redis_key = f"generated_image:{prompt}"
-    image = load_object_from_redis(redis_key)
-    if not image:
+    image_bytes = load_object_from_redis(redis_key)
+    if not image_bytes:
         try:
             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)
+            buffered = io.BytesIO()
+            image.save(buffered, format="JPEG")
+            image_bytes = buffered.getvalue()
+            save_object_to_redis(redis_key, image_bytes)
         except Exception as e:
             print(f"Failed to generate image: {e}")
             return None
-    return image
+    return image_bytes
 
-def edit_image_with_prompt(image, prompt, strength=0.75):
+def edit_image_with_prompt(image_bytes, prompt, strength=0.75):
     redis_key = f"edited_image:{prompt}:{strength}"
-    edited_image = load_object_from_redis(redis_key)
-    if not edited_image:
+    edited_image_bytes = load_object_from_redis(redis_key)
+    if not edited_image_bytes:
         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]
                 pbar.update(1)
-            save_object_to_redis(redis_key, edited_image)
+            buffered = io.BytesIO()
+            edited_image.save(buffered, format="JPEG")
+            edited_image_bytes = buffered.getvalue()
+            save_object_to_redis(redis_key, edited_image_bytes)
         except Exception as e:
             print(f"Failed to edit image: {e}")
             return None
-    return edited_image
+    return edited_image_bytes
 
 def generate_song(prompt, duration=10):
     redis_key = f"generated_song:{prompt}:{duration}"
@@ -128,8 +135,8 @@ def generate_text(prompt):
 
 def generate_flux_image(prompt):
     redis_key = f"generated_flux_image:{prompt}"
-    flux_image = load_object_from_redis(redis_key)
-    if not flux_image:
+    flux_image_bytes = load_object_from_redis(redis_key)
+    if not flux_image_bytes:
         try:
             with tqdm(total=1, desc="Generating FLUX image") as pbar:
                 flux_image = flux_pipeline(
@@ -140,11 +147,14 @@ def generate_flux_image(prompt):
                     generator=torch.Generator("cpu").manual_seed(0)
                 ).images[0]
                 pbar.update(1)
-            save_object_to_redis(redis_key, flux_image)
+            buffered = io.BytesIO()
+            flux_image.save(buffered, format="JPEG")
+            flux_image_bytes = buffered.getvalue()
+            save_object_to_redis(redis_key, flux_image_bytes)
         except Exception as e:
             print(f"Failed to generate flux image: {e}")
             return None
-    return flux_image
+    return flux_image_bytes
 
 def generate_code(prompt):
     redis_key = f"generated_code:{prompt}"
@@ -240,19 +250,19 @@ 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")
+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.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")
 
-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")
+gen_image_tab = gr.Interface(fn=generate_image, inputs=gr.Textbox(label="Prompt:"), outputs=gr.Image(type="pil"), title="Generate Image")
+edit_image_tab = gr.Interface(fn=edit_image_with_prompt, inputs=[gr.Image(type="pil", label="Image:"), gr.Textbox(label="Prompt:"), gr.Slider(0.1, 1.0, 0.75, step=0.05, label="Strength:")], outputs=gr.Image(type="pil"), title="Edit Image")
+generate_song_tab = gr.Interface(fn=generate_song, inputs=[gr.Textbox(label="Prompt:"), gr.Slider(5, 60, 10, step=1, label="Duration (s):")], outputs=gr.Audio(type="numpy"), title="Generate Songs")
+generate_text_tab = gr.Interface(fn=generate_text, inputs=gr.Textbox(label="Prompt:"), outputs=gr.Textbox(label="Generated Text:"), title="Generate Text")
+generate_flux_image_tab = gr.Interface(fn=generate_flux_image, inputs=gr.Textbox(label="Prompt:"), outputs=gr.Image(type="pil"), title="Generate FLUX Images")
+generate_code_tab = gr.Interface(fn=generate_code, inputs=gr.Textbox(label="Prompt:"), outputs=gr.Textbox(label="Generated Code:"), title="Generate Code")
+model_meta_llama_test_tab = gr.Interface(fn=test_model_meta_llama, inputs=gr.Textbox(label="Test Input:"), outputs=gr.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, generate_code_tab, model_meta_llama_test_tab],