Upload 7 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/sea.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,6 @@
 ---
 title: Noisescope
-emoji: 🦀
+emoji: 😎
 colorFrom: green
 colorTo: yellow
 sdk: gradio

app.py

@@ -0,0 +1,121 @@
+import uuid
+from pathlib import Path
+
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+from PIL import Image, ImageDraw, ImageFont
+from diffusers.schedulers import (
+    DDIMScheduler, DDPMScheduler,
+    PNDMScheduler
+)
+from gradio.themes import colors
+from torchvision.transforms.functional import to_pil_image, to_tensor
+from tqdm import tqdm
+
+SCHEDULER_CLASSES = {
+    "PNDMScheduler": PNDMScheduler,
+    "DDIMScheduler": DDIMScheduler,
+    "DDPMScheduler": DDPMScheduler,
+}
+
+
+@torch.inference_mode()
+def inference(
+        image: Image.Image,
+        scheduler_name: str,
+        num_timestep: int,
+        fps: int = 5,
+        seed: int = 1117,
+        progress=gr.Progress()
+):
+    progress(0, "Starting the process...")
+
+    # Set the output directory
+    out_dir = Path('out')
+    if out_dir.exists():
+        # Clean the directory
+        for file in out_dir.iterdir():
+            file.unlink()
+    else:
+        out_dir.mkdir()
+
+    model_id = 'CompVis/stable-diffusion-v1-4'
+
+    scheduler = SCHEDULER_CLASSES[scheduler_name].from_pretrained(model_id, subfolder='scheduler')
+    scheduler.set_timesteps(num_timestep)
+    timesteps = reversed(scheduler.timesteps)
+
+    # Resize image
+    img_width, img_height = image.size
+    image_tensor = to_tensor(image)
+
+    # Font info
+    # Adjust font size based on the image size
+    font_size = int(img_width / 30)
+    font = ImageFont.truetype("assets/JetBrainsMono-Bold.ttf", font_size)
+
+    # Prepare and Add noise
+    generator = torch.Generator().manual_seed(seed)
+    noise = torch.randn(image_tensor.shape, generator=generator)
+
+    # Define the codec and create VideoWriter object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Be sure to use lower case
+    # fps = 5
+
+    # Create an in-memory buffer
+    out_path = f'{str(out_dir)}/{uuid.uuid4()}.mp4'
+    out = cv2.VideoWriter(out_path, fourcc, fps, (img_width, img_height))
+
+    for i, timestep in enumerate(tqdm(timesteps, desc="Writing video frames")):
+        noised_image = scheduler.add_noise(image_tensor, noise, timestep)
+        noised_image = to_pil_image(noised_image)
+
+        # Add caption to the image
+        draw = ImageDraw.Draw(noised_image)
+        caption = f"Iteration: {i + 1}, Timestep: {timestep.item()}"
+
+        cap_x = img_width / 4
+        cap_y = 10 + (font_size * 2)
+
+        draw.text(xy=(cap_x, cap_y), text=caption, fill=(255, 255, 255), font=font)
+
+        noised_image = cv2.cvtColor(np.array(noised_image), cv2.COLOR_RGB2BGR)
+        out.write(noised_image)  # Write out frame to video
+        progress(1 / num_timestep * (i + 1), f"Writing frame {i + 1} of {num_timestep}...")
+
+    out.release()  # Release the video writer
+    progress(1, "Video writing completed.")
+    return out_path
+
+
+if __name__ == '__main__':
+    iface = gr.Interface(
+        fn=inference,
+        inputs=[
+            gr.Image(type="pil", label="Input Image"),
+            gr.Dropdown(
+                choices=[name for name in SCHEDULER_CLASSES.keys()],
+                label="Scheduler",
+                value="PNDMScheduler",
+            ),
+            gr.Slider(minimum=10, maximum=990, value=50, step=10, label="Total Timesteps"),
+            gr.Slider(minimum=1, maximum=60, value=5, step=1, label="fps"),
+            gr.Number(value=1117, label="Seed"),
+        ],
+        outputs=[
+            gr.PlayableVideo(label="Process of adding noise")
+        ],
+        examples=[
+            [Image.open("assets/sea.jpg"), "PNDMScheduler", 999, 1117],
+            [Image.open("assets/ckwb.jpg"), "PNDMScheduler", 30, 1117],
+            [Image.open("assets/dog.jpg"), "PNDMScheduler", 50, 1117],
+        ],
+        cache_examples=True,
+        title="NoiseScope 😎",
+        description="Visualize what amount of noise is added at each time step in the forward process of the diffusion model 🧨",
+        theme=gr.Theme(primary_hue=colors.green, secondary_hue=colors.yellow),
+        allow_flagging="never",
+    )
+    iface.queue().launch()

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+diffusers
+transformers
+safetensors
+accelerate
+opencv-python