Update app.py

app.py

@@ -1,5 +1,8 @@
 import torch
 import gradio as gr
+import numpy as np
+import random
+
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionInstructPix2PixPipeline,
@@ -7,59 +10,65 @@ 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
+# Force CPU mode
+device = "cpu"
+dtype = torch.float32
 MAX_SEED = np.iinfo(np.int32).max
 
-# Model cache
+# Global pipeline holders
 TXT2IMG_PIPE = None
 IMG2IMG_PIPE = None
 TXT2VID_PIPE = None
 IMG2VID_PIPE = None
 
+# Helper to load models
 def make_pipe(cls, model_id, **kwargs):
     pipe = cls.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
-    pipe.enable_model_cpu_offload()
+    pipe.to(device)
     return pipe
 
-# Functions
+# Text → Image
 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)
+        TXT2IMG_PIPE = make_pipe(StableDiffusionPipeline, "stabilityai/stable-diffusion-2-1-base")
     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, seed, randomize_seed):
     global IMG2IMG_PIPE
     if IMG2IMG_PIPE is None:
-        IMG2IMG_PIPE = make_pipe(StableDiffusionInstructPix2PixPipeline, "timbrooks/instruct-pix2pix").to(device)
+        IMG2IMG_PIPE = make_pipe(StableDiffusionInstructPix2PixPipeline, "timbrooks/instruct-pix2pix")
     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, 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)
+        TXT2VID_PIPE = make_pipe(WanPipeline, "Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
     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, 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"
+        )
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.manual_seed(seed)
@@ -67,65 +76,61 @@ def generate_video_from_image(image, seed, randomize_seed):
     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
 
-# UI
-with gr.Blocks(css="footer {display:none !important}") as demo:
-    gr.Markdown("# 🧠 AI Playground – Multi-Mode Generator")
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 AI Playground – Text & Image → Image & Video")
 
     with gr.Tabs():
         # Text → Image
         with gr.Tab("Text → Image"):
-            with gr.Row():
-                prompt_txt = gr.Textbox(label="Prompt")
-                generate_btn = gr.Button("Generate")
+            prompt_txt = gr.Textbox(label="Prompt")
+            btn_txt2img = 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],
+            rand_txt = gr.Checkbox(label="Randomize seed", value=True)
+            btn_txt2img.click(
+                generate_image_from_text,
+                inputs=[prompt_txt, seed_txt, rand_txt],
                 outputs=[result_img, seed_txt]
             )
 
         # 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")
+            image_in = gr.Image(label="Input Image")
+            prompt_img = gr.Textbox(label="Edit Prompt")
+            btn_img2img = 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],
+            rand_img = gr.Checkbox(label="Randomize seed", value=True)
+            btn_img2img.click(
+                generate_image_from_image_and_prompt,
+                inputs=[image_in, prompt_img, seed_img, rand_img],
                 outputs=[result_img2, seed_img]
             )
 
         # Text → Video
         with gr.Tab("Text → Video"):
-            with gr.Row():
-                prompt_vid = gr.Textbox(label="Prompt")
-                generate_btn3 = gr.Button("Generate")
+            prompt_vid = gr.Textbox(label="Prompt")
+            btn_txt2vid = 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],
+            rand_vid = gr.Checkbox(label="Randomize seed", value=True)
+            btn_txt2vid.click(
+                generate_video_from_text,
+                inputs=[prompt_vid, seed_vid, rand_vid],
                 outputs=[result_vid, seed_vid]
             )
 
         # Image → Video
         with gr.Tab("Image → Video"):
-            with gr.Row():
-                image_in_vid = gr.Image(label="Input Image")
-                generate_btn4 = gr.Button("Animate")
+            image_vid = gr.Image(label="Input Image")
+            btn_img2vid = 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],
+            rand_vid2 = gr.Checkbox(label="Randomize seed", value=True)
+            btn_img2vid.click(
+                generate_video_from_image,
+                inputs=[image_vid, seed_vid2, rand_vid2],
                 outputs=[result_vid2, seed_vid2]
             )