Create app.py

app.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+import os
+import inspect
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import numpy as np
+from torch import autocast
+import cv2
+import gradio as gr
+
+# -----------------------------------------------------------------------------
+# 1. REQUIREMENTS & SETUP
+# -----------------------------------------------------------------------------
+# To set up the environment for this script, create a file named 'requirements.txt'
+# with the following content and run 'pip install -r requirements.txt':
+#
+# torch>=2.0.0
+# torchvision>=0.15.1
+# diffusers>=0.20.2
+# transformers>=4.30.2
+# accelerate>=0.21.0
+# gradio>=3.36.1
+# opencv-python-headless>=4.8.0.74
+# -----------------------------------------------------------------------------
+
+# --- Automatic Device Detection ---
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+print("-------------------------------------------------")
+print(f"INFO: Using device: {torch_device.upper()}")
+if torch_device == "cpu":
+    print("WARNING: CUDA (GPU) not detected. The script will run on the CPU.")
+    print("         This will be extremely slow. For better performance,")
+    print("         please ensure you have an NVIDIA GPU and the correct")
+    print("         PyTorch version with CUDA support installed.")
+print("-------------------------------------------------")
+
+
+# --- Load the Model ---
+print("Loading Stable Diffusion model... This may take a moment.")
+try:
+    # Load the pipeline and move it to the detected device
+    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
+    pipe.to(torch_device)
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    print("Please check your internet connection and ensure the model name is correct.")
+    exit()
+
+# -----------------------------------------------------------------------------
+# Helper Functions (slerp, diffuse)
+# -----------------------------------------------------------------------------
+
+@torch.no_grad()
+def diffuse(
+        pipe, cond_embeddings, cond_latents, num_inference_steps, guidance_scale, eta
+    ):
+    # This function remains the same, as it gets the device from the input tensors
+    device = cond_latents.get_device()
+    max_length = cond_embeddings.shape[1]
+    uncond_input = pipe.tokenizer([""], padding="max_length", max_length=max_length, return_tensors="pt")
+    uncond_embeddings = pipe.text_encoder(uncond_input.input_ids.to(device))[0]
+    text_embeddings = torch.cat([uncond_embeddings, cond_embeddings])
+
+    if "LMS" in pipe.scheduler.__class__.__name__:
+        cond_latents = cond_latents * pipe.scheduler.sigmas[0]
+
+    accepts_offset = "offset" in set(inspect.signature(pipe.scheduler.set_timesteps).parameters.keys())
+    extra_set_kwargs = {}
+    if accepts_offset:
+        extra_set_kwargs["offset"] = 1
+    pipe.scheduler.set_timesteps(num_inference_steps, **extra_set_kwargs)
+
+    accepts_eta = "eta" in set(inspect.signature(pipe.scheduler.step).parameters.keys())
+    extra_step_kwargs = {}
+    if accepts_eta:
+        extra_step_kwargs["eta"] = eta
+
+    for i, t in enumerate(pipe.scheduler.timesteps):
+        latent_model_input = torch.cat([cond_latents] * 2)
+        if "LMS" in pipe.scheduler.__class__.__name__:
+            sigma = pipe.scheduler.sigmas[i]
+            latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)
+
+        # predict the noise residual
+        noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings)["sample"]
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        cond_latents = pipe.scheduler.step(noise_pred, t, cond_latents, **extra_step_kwargs)["prev_sample"]
+
+    cond_latents = 1 / 0.18215 * cond_latents
+    image = pipe.vae.decode(cond_latents).sample
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = image.cpu().permute(0, 2, 3, 1).numpy()
+    image = (image[0] * 255).astype(np.uint8)
+    return image
+
+def slerp(t, v0, v1, DOT_THRESHOLD=0.9995):
+    # This function is device-agnostic
+    inputs_are_torch = isinstance(v0, torch.Tensor)
+    if inputs_are_torch:
+        input_device = v0.device
+        v0 = v0.cpu().numpy()
+        v1 = v1.cpu().numpy()
+    dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+    if np.abs(dot) > DOT_THRESHOLD:
+        v2 = (1 - t) * v0 + t * v1
+    else:
+        theta_0 = np.arccos(dot)
+        sin_theta_0 = np.sin(theta_0)
+        theta_t = theta_0 * t
+        sin_theta_t = np.sin(theta_t)
+        s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+        s1 = sin_theta_t / sin_theta_0
+        v2 = s0 * v0 + s1 * v1
+    if inputs_are_torch:
+        v2 = torch.from_numpy(v2).to(input_device)
+    return v2
+
+# -----------------------------------------------------------------------------
+# Main Generator Function for Gradio
+# -----------------------------------------------------------------------------
+def generate_dream_video(
+    prompt_1, prompt_2, seed_1, seed_2,
+    width, height, num_steps, guidance_scale,
+    num_inference_steps, eta, name
+):
+    # --- 1. SETUP ---
+    yield {
+        status_text: "Status: Preparing prompts and latents...",
+        live_frame: None,
+        output_video: None,
+    }
+    prompts = [prompt_1, prompt_2]
+    seeds = [int(seed_1), int(seed_2)]
+    rootdir = './dreams'
+    outdir = os.path.join(rootdir, name)
+    os.makedirs(outdir, exist_ok=True)
+    
+    # --- 2. EMBEDDINGS AND LATENTS ---
+    prompt_embeddings = []
+    for prompt in prompts:
+        text_input = pipe.tokenizer(prompt, padding="max_length", max_length=pipe.tokenizer.model_max_length, truncation=True, return_tensors="pt")
+        # Move input_ids to the correct device before text encoding
+        with torch.no_grad():
+            embed = pipe.text_encoder(text_input.input_ids.to(torch_device))[0]
+        prompt_embeddings.append(embed)
+
+    prompt_embedding_a, prompt_embedding_b = prompt_embeddings
+    
+    # Use a device-specific generator for reproducibility
+    generator_a = torch.Generator(device=torch_device).manual_seed(seeds[0])
+    generator_b = torch.Generator(device=torch_device).manual_seed(seeds[1])
+
+    init_a = torch.randn((1, pipe.unet.config.in_channels, height // 8, width // 8), device=torch_device, generator=generator_a)
+    init_b = torch.randn((1, pipe.unet.config.in_channels, height // 8, width // 8), device=torch_device, generator=generator_b)
+
+    # --- 3. GENERATION LOOP ---
+    frame_paths = []
+    for i, t in enumerate(np.linspace(0, 1, num_steps)):
+        yield {
+            status_text: f"Status: Generating frame {i + 1} of {num_steps} on {torch_device.upper()}...",
+            live_frame: None,
+            output_video: None,
+        }
+        
+        cond_embedding = slerp(float(t), prompt_embedding_a, prompt_embedding_b)
+        init = slerp(float(t), init_a, init_b)
+
+        # Use autocast only if on CUDA
+        with autocast(torch_device) if torch_device == "cuda" else open(os.devnull, 'w') as f:
+            image = diffuse(pipe, cond_embedding, init, num_inference_steps, guidance_scale, eta)
+
+        im = Image.fromarray(image)
+        outpath = os.path.join(outdir, f'frame{i:06d}.jpg')
+        im.save(outpath)
+        frame_paths.append(outpath)
+        
+        yield { live_frame: im }
+
+    # --- 4. VIDEO COMPILATION ---
+    yield { status_text: "Status: Compiling video from frames..." }
+    
+    video_path = os.path.join(outdir, f"{name}.mp4")
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(video_path, fourcc, 15, (width, height))
+    for frame_path in frame_paths:
+        frame = cv2.imread(frame_path)
+        video_writer.write(frame)
+    video_writer.release()
+    
+    print(f"Video saved to {video_path}")
+    yield {
+        status_text: f"Status: Done! Video saved to {video_path}",
+        output_video: video_path
+    }
+
+# -----------------------------------------------------------------------------
+# Gradio UI (Unchanged)
+# -----------------------------------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft(), css="footer {display: none !important}") as demo:
+    gr.Markdown("# 🎥 Stable Diffusion Video Interpolation")
+    gr.Markdown("Create smooth transition videos between two concepts. Configure the prompts and settings below, then click Generate.")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Accordion("1. Core Prompts & Seeds", open=True):
+                prompt_1 = gr.Textbox(lines=2, label="Starting Prompt", value="ultrarealistic steam punk neural network machine in the shape of a brain, placed on a pedestal, covered with neurons made of gears.")
+                seed_1 = gr.Number(label="Seed 1", value=243, precision=0, info="A specific number to control the starting noise pattern.")
+                prompt_2 = gr.Textbox(lines=2, label="Ending Prompt", value="A bioluminescent, glowing jellyfish floating in a dark, deep abyss, surrounded by sparkling plankton.")
+                seed_2 = gr.Number(label="Seed 2", value=523, precision=0, info="A specific number to control the ending noise pattern.")
+                name = gr.Textbox(label="Output File Name", value="my_dream_video", info="The name for the output folder and .mp4 file.")
+            
+            with gr.Accordion("2. Generation Parameters", open=True):
+                with gr.Row():
+                    width = gr.Slider(label="Width", minimum=256, maximum=1024, value=512, step=64)
+                    height = gr.Slider(label="Height", minimum=256, maximum=1024, value=512, step=64)
+                num_steps = gr.Slider(label="Interpolation Frames", minimum=10, maximum=500, value=120, step=1, info="How many frames the final video will have. More frames = smoother video.")
+
+            with gr.Accordion("3. Advanced Diffusion Settings", open=False):
+                 num_inference_steps = gr.Slider(label="Inference Steps per Frame", minimum=10, maximum=100, value=40, step=1, info="More steps can improve quality but will be much slower.")
+                 guidance_scale = gr.Slider(label="Guidance Scale (CFG)", minimum=1, maximum=20, value=7.5, step=0.5, info="How strongly the prompt guides the image generation.")
+                 eta = gr.Slider(label="ETA (for DDIM Scheduler)", minimum=0.0, maximum=1.0, value=0.0, step=0.1, info="A parameter for noise scheduling. 0.0 is deterministic.")
+
+            run_button = gr.Button("Generate Video", variant="primary")
+
+        with gr.Column(scale=3):
+            status_text = gr.Textbox(label="Status", value="Ready", interactive=False)
+            live_frame = gr.Image(label="Live Preview", type="pil")
+            output_video = gr.Video(label="Final Video")
+
+    run_button.click(
+        fn=generate_dream_video,
+        inputs=[
+            prompt_1, prompt_2, seed_1, seed_2,
+            width, height, num_steps, guidance_scale,
+            num_inference_steps, eta, name
+        ],
+        outputs=[status_text, live_frame, output_video]
+    )
+
+# --- Launch the App ---
+if __name__ == "__main__":
+    demo.launch(share=True, debug=True)