Update app.py

app.py

@@ -1,80 +1,68 @@
 from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
-from utils import write_video, dummy
+from utils import write_video, dummy, preprocess_image, preprocess_mask_image
 from PIL import Image
-import numpy as np
+import gradio as gr
+import torch
 import os
 os.environ["CUDA_VISIBLE_DEVICES"]="0"
-import torch
-import gradio as gr
 
 
-orig_prompt = "Create a relaxing atmosphere with the use of plants and other natural elements. Such as a hanging terrarium or a wall-mounted planter. Include plenty of storage options to keep the space organized and clutter-free. Consider adding a vanity with double sinks and plenty of drawers and cabinets. As well as a wall mounted medicine and towel storage."
+orig_prompt = "Ancient underground architectural ruins of Hong Kong in a flooded apocalypse landscape of dead skyscrapers"
 orig_negative_prompt = "lurry, bad art, blurred, text, watermark"
+model_list = ["stabilityai/stable-diffusion-2-inpainting", "runwayml/stable-diffusion-inpainting"]
 
 def stable_diffusion_zoom_out(
   repo_id,
   original_prompt,
   negative_prompt,
-  steps,
+  step_size,
   num_frames,
-  fps
-):
+  fps,
+  num_inference_steps
+    ):
     
-  pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, revision="fp16")
-  pipe.set_use_memory_efficient_attention_xformers(True)
-  pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-  pipe = pipe.to("cuda")
-  pipe.safety_checker = dummy
-
-  current_image = Image.new(mode="RGBA", size=(512,512))
-  mask_image = np.array(current_image)[:,:,3] # assume image has alpha mask (use .mode to check for "RGBA")
-  mask_image = Image.fromarray(255-mask_image).convert("RGB")
-  current_image = current_image.convert("RGB")
+    pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16)
+    pipe.set_use_memory_efficient_attention_xformers(True)
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe = pipe.to("cuda")
+    pipe.safety_checker = dummy
 
-  num_images = 1
-  prompt = [original_prompt] * num_images
-  negative_prompt = [negative_prompt] * num_images
+    new_image = Image.new(mode="RGBA", size=(512,512)) 
+    current_image, mask_image = preprocess_mask_image(new_image)
+    
+    current_image = pipe(prompt=[original_prompt], negative_prompt=[negative_prompt], image=current_image, mask_image=mask_image, num_inference_steps=num_inference_steps).images[0]
 
-  images = pipe(prompt=prompt, negative_prompt=negative_prompt, image=current_image, mask_image=mask_image, num_inference_steps=25)[0]
-  current_image = images[0]
+    
+    all_frames = []
+    all_frames.append(current_image)
 
-  all_frames = []
-  all_frames.append(current_image)
+    for i in range(num_frames):
+        prev_image = preprocess_image(current_image, step_size, 512)
+        current_image = prev_image
+        current_image, mask_image = preprocess_mask_image(current_image)
+        current_image = pipe(prompt=[original_prompt], negative_prompt=[negative_prompt], image=current_image, mask_image=mask_image, num_inference_steps=num_inference_steps).images[0]
+        
+        current_image.paste(prev_image, mask=prev_image)
+        all_frames.append(current_image)
 
-  for i in range(num_frames):
-    next_image = np.array(current_image.convert("RGBA"))*0
-    prev_image = current_image.resize((512-2*steps,512-2*steps))
-    prev_image = prev_image.convert("RGBA")
-    prev_image = np.array(prev_image)
-    next_image[:, :, 3] = 1
-    next_image[steps:512-steps,steps:512-steps,:] = prev_image
-    prev_image = Image.fromarray(next_image)
-    current_image = prev_image
-    mask_image = np.array(current_image)[:,:,3] # assume image has alpha mask (use .mode to check for "RGBA")
-    mask_image = Image.fromarray(255-mask_image).convert("RGB")
-    current_image = current_image.convert("RGB")
-    images = pipe(prompt=prompt, negative_prompt=negative_prompt, image=current_image, mask_image=mask_image, num_inference_steps=25)[0]
-    current_image = images[0]
-    current_image.paste(prev_image, mask=prev_image)
-    all_frames.append(current_image)
-    
-  save_path = "infinite_zoom_out.mp4"  
-  write_video(save_path, all_frames, fps=fps)
-  return save_path
+    save_path = "output.mp4"  
+    write_video(save_path, all_frames, fps=fps)
+    return save_path
 
 
 inputs = [
-  gr.inputs.Textbox(lines=1, default="stabilityai/stable-diffusion-2-inpainting", label="Model ID"),
+  gr.Dropdown(model_list, value=model_list[0], label="Model"),
   gr.inputs.Textbox(lines=5, default=orig_prompt, label="Prompt"),
   gr.inputs.Textbox(lines=1, default=orig_negative_prompt, label="Negative Prompt"),
-  gr.inputs.Slider(minimum=1, maximum=64, default=32, label="Steps"),
-  gr.inputs.Slider(minimum=1, maximum=500, default=10, step=10, label="Frames"),
-  gr.inputs.Slider(minimum=1, maximum=100, default=16, step=1, label="FPS")
+  gr.inputs.Slider(minimum=1, maximum=120, default=25, step=5, label="Steps"),
+  gr.inputs.Slider(minimum=1, maximum=100, default=10, step=5, label="Frames"),
+  gr.inputs.Slider(minimum=1, maximum=100, default=16, step=1, label="FPS"),
+  gr.inputs.Slider(minimum=1, maximum=100, default=15, step=1, label="Inference Steps")
 ]
 
 output = gr.outputs.Video()
 examples = [
-  ["stabilityai/stable-diffusion-2-inpainting", orig_prompt, orig_negative_prompt, 32, 50, 16]
+  ["stabilityai/stable-diffusion-2-inpainting", orig_prompt, orig_negative_prompt, 25, 10, 16, 15],
 ]
 
 title = "Stable Diffusion Infinite Zoom Out"
@@ -84,7 +72,6 @@ description = """<p>For faster inference without waiting in queue, you may dupli
 <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
 <p/>"""
 
-
 demo_app = gr.Interface(
     fn=stable_diffusion_zoom_out,
     description=description,