Update app.py

app.py

@@ -7,95 +7,127 @@ from diffusers import (
     WanPipeline,
 )
 from diffusers.utils import export_to_video, load_image
+import random
+import numpy as np
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16 if device == "cuda" else torch.float32
+MAX_SEED = np.iinfo(np.int32).max
 
-# Pipeline factory
-def make_pipe(cls, model_id, **kwargs):
-    pipe = cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
-    pipe.enable_model_cpu_offload()
-    return pipe
-
-# Global model caches
+# Model cache
 TXT2IMG_PIPE = None
 IMG2IMG_PIPE = None
 TXT2VID_PIPE = None
 IMG2VID_PIPE = None
 
-# Text → Image
-def generate_image_from_text(prompt):
+def make_pipe(cls, model_id, **kwargs):
+    pipe = cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
+    pipe.enable_model_cpu_offload()
+    return pipe
+
+# Functions
+def generate_image_from_text(prompt, seed, randomize_seed):
     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]
+        TXT2IMG_PIPE = make_pipe(StableDiffusionPipeline, "stabilityai/stable-diffusion-2-1-base").to(device)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.manual_seed(seed)
+    image = TXT2IMG_PIPE(prompt=prompt, num_inference_steps=20, generator=generator).images[0]
+    return image, seed
 
-# Image → Image
-def generate_image_from_image_and_prompt(image, prompt):
+def generate_image_from_image_and_prompt(image, prompt, seed, randomize_seed):
     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]
+        IMG2IMG_PIPE = make_pipe(StableDiffusionInstructPix2PixPipeline, "timbrooks/instruct-pix2pix").to(device)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.manual_seed(seed)
+    out = IMG2IMG_PIPE(prompt=prompt, image=image, num_inference_steps=8, generator=generator)
+    return out.images[0], seed
 
-# Text → Video
-def generate_video_from_text(prompt):
+def generate_video_from_text(prompt, seed, randomize_seed):
     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, "/tmp/wan_video.mp4", fps=8)
+        TXT2VID_PIPE = make_pipe(WanPipeline, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers").to(device)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.manual_seed(seed)
+    frames = TXT2VID_PIPE(prompt=prompt, num_frames=12, generator=generator).frames[0]
+    return export_to_video(frames, "/tmp/wan_video.mp4", fps=8), seed
 
-# Image → Video
-def generate_video_from_image(image):
+def generate_video_from_image(image, seed, randomize_seed):
     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)
+        IMG2VID_PIPE = make_pipe(StableVideoDiffusionPipeline, "stabilityai/stable-video-diffusion-img2vid-xt", variant="fp16" if dtype == torch.float16 else None).to(device)
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.manual_seed(seed)
     image = load_image(image).resize((512, 288))
-    frames = IMG2VID_PIPE(image, num_inference_steps=16).frames[0]
-    return export_to_video(frames, "/tmp/svd_video.mp4", fps=8)
+    frames = IMG2VID_PIPE(image=image, num_inference_steps=16, generator=generator).frames[0]
+    return export_to_video(frames, "/tmp/svd_video.mp4", fps=8), seed
 
-# Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Lightweight Any‑to‑Any AI Playground")
+# UI
+with gr.Blocks(css="footer {display:none !important}") as demo:
+    gr.Markdown("# 🧠 AI Playground – Multi-Mode Generator")
 
-    with gr.Tab("Text → Image"):
-        text_prompt = gr.Textbox(label="Prompt")
-        output_image = gr.Image(label="Generated Image")
-        text2img_button = gr.Button("Generate")
-        text2img_button.click(generate_image_from_text, inputs=text_prompt, outputs=output_image)
+    with gr.Tabs():
+        # Text → Image
+        with gr.Tab("Text → Image"):
+            with gr.Row():
+                prompt_txt = gr.Textbox(label="Prompt")
+                generate_btn = gr.Button("Generate")
+            result_img = gr.Image()
+            seed_txt = gr.Slider(0, MAX_SEED, value=42, label="Seed")
+            rand_seed_txt = gr.Checkbox(label="Randomize seed", value=True)
+            generate_btn.click(
+                fn=generate_image_from_text,
+                inputs=[prompt_txt, seed_txt, rand_seed_txt],
+                outputs=[result_img, seed_txt]
+            )
 
-    with gr.Tab("Image → Image"):
-        input_image = gr.Image(label="Input Image")
-        edit_prompt = gr.Textbox(label="Edit Prompt")
-        edited_image = gr.Image(label="Edited Image")
-        img2img_button = gr.Button("Generate")
-        img2img_button.click(generate_image_from_image_and_prompt, inputs=[input_image, edit_prompt], outputs=edited_image)
+        # Image → Image
+        with gr.Tab("Image → Image"):
+            with gr.Row():
+                image_in = gr.Image(label="Input Image")
+                prompt_img = gr.Textbox(label="Edit Prompt")
+                generate_btn2 = gr.Button("Generate")
+            result_img2 = gr.Image()
+            seed_img = gr.Slider(0, MAX_SEED, value=123, label="Seed")
+            rand_seed_img = gr.Checkbox(label="Randomize seed", value=True)
+            generate_btn2.click(
+                fn=generate_image_from_image_and_prompt,
+                inputs=[image_in, prompt_img, seed_img, rand_seed_img],
+                outputs=[result_img2, seed_img]
+            )
 
-    with gr.Tab("Text → Video"):
-        video_prompt = gr.Textbox(label="Prompt")
-        video_output = gr.Video(label="Generated Video")
-        txt2vid_button = gr.Button("Generate")
-        txt2vid_button.click(generate_video_from_text, inputs=video_prompt, outputs=video_output)
+        # Text → Video
+        with gr.Tab("Text → Video"):
+            with gr.Row():
+                prompt_vid = gr.Textbox(label="Prompt")
+                generate_btn3 = gr.Button("Generate")
+            result_vid = gr.Video()
+            seed_vid = gr.Slider(0, MAX_SEED, value=555, label="Seed")
+            rand_seed_vid = gr.Checkbox(label="Randomize seed", value=True)
+            generate_btn3.click(
+                fn=generate_video_from_text,
+                inputs=[prompt_vid, seed_vid, rand_seed_vid],
+                outputs=[result_vid, seed_vid]
+            )
 
-    with gr.Tab("Image → Video"):
-        video_input_img = gr.Image(label="Input Image")
-        anim_video_output = gr.Video(label="Animated Video")
-        img2vid_button = gr.Button("Animate")
-        img2vid_button.click(generate_video_from_image, inputs=video_input_img, outputs=anim_video_output)
+        # Image → Video
+        with gr.Tab("Image → Video"):
+            with gr.Row():
+                image_in_vid = gr.Image(label="Input Image")
+                generate_btn4 = gr.Button("Animate")
+            result_vid2 = gr.Video()
+            seed_vid2 = gr.Slider(0, MAX_SEED, value=999, label="Seed")
+            rand_seed_vid2 = gr.Checkbox(label="Randomize seed", value=True)
+            generate_btn4.click(
+                fn=generate_video_from_image,
+                inputs=[image_in_vid, seed_vid2, rand_seed_vid2],
+                outputs=[result_vid2, seed_vid2]
+            )
 
 demo.queue()
 demo.launch(show_error=True)