Update app.py

app.py

@@ -8,86 +8,86 @@ from diffusers import (
 )
 from diffusers.utils import export_to_video, load_image
 
-# Set dtype and device
-dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+# Detect device & dtype
 device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
 
-# -------- Text to Image: Stable Diffusion --------
-txt2img_pipe = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1-base", torch_dtype=dtype
-)
-txt2img_pipe.to(device)
-
-def generate_image_from_text(prompt):
-    image = txt2img_pipe(prompt, num_inference_steps=30).images[0]
-    return image
+# Factory to load & offload a pipeline
+def make_pipe(cls, model_id, **kwargs):
+    pipe = cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
+    # Enables CPU offload of model parts not in use
+    pipe.enable_model_cpu_offload()
+    return pipe
 
+# Hold pipelines in globals but don’t load yet  
+TXT2IMG_PIPE = None
+IMG2IMG_PIPE = None
+TXT2VID_PIPE = None
+IMG2VID_PIPE = None
 
-# -------- Image to Image: Instruct Pix2Pix --------
-pix2pix_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-    "timbrooks/instruct-pix2pix", torch_dtype=dtype
-)
-pix2pix_pipe.to(device)
+def generate_image_from_text(prompt):
+    global TXT2IMG_PIPE
+    if TXT2IMG_PIPE is None:
+        TXT2IMG_PIPE = make_pipe(
+            StableDiffusionPipeline, 
+            "stabilityai/stable-diffusion-2-1-base"
+        ).to(device)
+    return TXT2IMG_PIPE(prompt, num_inference_steps=20).images[0]
 
 def generate_image_from_image_and_prompt(image, prompt):
-    result = pix2pix_pipe(prompt=prompt, image=image, num_inference_steps=10)
-    return result.images[0]
-
-
-# -------- Text to Video: Wan T2V --------
-wan_pipe = WanPipeline.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", torch_dtype=torch.bfloat16
-)
-wan_pipe.to(device)
+    global IMG2IMG_PIPE
+    if IMG2IMG_PIPE is None:
+        IMG2IMG_PIPE = make_pipe(
+            StableDiffusionInstructPix2PixPipeline,
+            "timbrooks/instruct-pix2pix"
+        ).to(device)
+    out = IMG2IMG_PIPE(prompt=prompt, image=image, num_inference_steps=8)
+    return out.images[0]
 
 def generate_video_from_text(prompt):
-    frames = wan_pipe(prompt=prompt, num_frames=16).frames[0]
-    video_path = export_to_video(frames, "wan_video.mp4", fps=8)
-    return video_path
-
-
-# -------- Image to Video: Stable Video Diffusion --------
-svd_pipe = StableVideoDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-video-diffusion-img2vid-xt",
-    torch_dtype=dtype,
-    variant="fp16" if dtype == torch.float16 else None,
-)
-svd_pipe.to(device)
+    global TXT2VID_PIPE
+    if TXT2VID_PIPE is None:
+        TXT2VID_PIPE = make_pipe(
+            WanPipeline,
+            "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+        ).to(device)
+    frames = TXT2VID_PIPE(prompt=prompt, num_frames=12).frames[0]
+    return export_to_video(frames, "wan_video.mp4", fps=8)
 
 def generate_video_from_image(image):
-    image = image.resize((1024, 576))
-    frames = svd_pipe(image, num_inference_steps=25).frames[0]
-    video_path = export_to_video(frames, "svd_video.mp4", fps=8)
-    return video_path
-
+    global IMG2VID_PIPE
+    if IMG2VID_PIPE is None:
+        IMG2VID_PIPE = make_pipe(
+            StableVideoDiffusionPipeline,
+            "stabilityai/stable-video-diffusion-img2vid-xt",
+            variant="fp16" if dtype==torch.float16 else None
+        ).to(device)
+    image = load_image(image).resize((512, 288))
+    frames = IMG2VID_PIPE(image, num_inference_steps=16).frames[0]
+    return export_to_video(frames, "svd_video.mp4", fps=8)
 
-# -------- Gradio Interface --------
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Multimodal Any-to-Any AI Playground")
+    gr.Markdown("# 🧠 Lightweight Any‑to‑Any AI Playground")
 
     with gr.Tab("Text → Image"):
-        prompt = gr.Textbox(label="Prompt")
-        output_image = gr.Image()
-        btn1 = gr.Button("Generate")
-        btn1.click(fn=generate_image_from_text, inputs=prompt, outputs=output_image)
+        inp = gr.Textbox(label="Prompt")
+        out = gr.Image()
+        gr.Button("Generate").click(generate_image_from_text, inp, out)
 
     with gr.Tab("Image → Image"):
-        in_image = gr.Image(label="Input Image")
-        edit_prompt = gr.Textbox(label="Edit Prompt")
-        out_image = gr.Image()
-        btn2 = gr.Button("Generate")
-        btn2.click(fn=generate_image_from_image_and_prompt, inputs=[in_image, edit_prompt], outputs=out_image)
+        img = gr.Image(label="Input Image")
+        prm = gr.Textbox(label="Edit Prompt")
+        out2 = gr.Image()
+        gr.Button("Generate").click(generate_image_from_image_and_prompt, [img, prm], out2)
 
     with gr.Tab("Text → Video"):
-        vid_prompt = gr.Textbox(label="Prompt")
-        output_vid = gr.Video()
-        btn3 = gr.Button("Generate")
-        btn3.click(fn=generate_video_from_text, inputs=vid_prompt, outputs=output_vid)
+        inp2 = gr.Textbox(label="Prompt")
+        out_vid = gr.Video()
+        gr.Button("Generate").click(generate_video_from_text, inp2, out_vid)
 
     with gr.Tab("Image → Video"):
-        img_input = gr.Image(label="Input Image")
-        vid_out = gr.Video()
-        btn4 = gr.Button("Animate")
-        btn4.click(fn=generate_video_from_image, inputs=img_input, outputs=vid_out)
+        img2 = gr.Image(label="Input Image")
+        out_vid2 = gr.Video()
+        gr.Button("Animate").click(generate_video_from_image, img2, out_vid2)
 
 demo.launch()