# -*- 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 # torchvision # diffusers # transformers # accelerate # gradio # opencv-python-headless # ----------------------------------------------------------------------------- # --- 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, device ): # The 'device' is now passed explicitly to this function max_length = cond_embeddings.shape[1] uncond_input = pipe.tokenizer([""], padding="max_length", max_length=max_length, return_tensors="pt") # Use the passed 'device' variable for all tensor placement 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") 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 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: # Pass the torch_device explicitly to the diffuse function image = diffuse(pipe, cond_embedding, init, num_inference_steps, guidance_scale, eta, torch_device) 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)