first commit

app.py

@@ -0,0 +1,163 @@
+from PIL import Image
+import numpy as np
+import gradio as gr
+import spaces
+import torch
+from tqdm import tqdm
+
+from controlnet import QRControlNet
+from game_of_life import GameOfLife
+from utils import resize_image, generate_image_from_grid
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+controlnet = QRControlNet(device=device)
+
+
+def generate_all_images(
+    gol_grids: list[np.array],
+    source_image: Image,
+    num_inference_steps: int,
+    controlnet_conditioning_scale: float,
+    strength: float,
+    prompt: str,
+    negative_prompt: str,
+    seed: int,
+    guidance_scale: float,
+    img_size: int,
+):
+
+    controlnet_conditioning_scale = float(controlnet_conditioning_scale)
+    source_image = resize_image(source_image, resolution=img_size)
+    images = []
+    for grid in tqdm(gol_grids):
+
+        grid_inverse = 1 - grid  # invert the grid for controlnet
+        grid_inverse_image = generate_image_from_grid(grid_inverse, img_size=img_size)
+
+        image = controlnet.generate_image(
+            source_image=source_image,
+            control_image=grid_inverse_image,
+            num_inference_steps=num_inference_steps,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            strength=strength,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            seed=seed,
+            guidance_scale=guidance_scale,
+            img_size=img_size,
+        )
+        images.append(image)
+
+    return images
+
+
+def make_gif(images: list[Image.Image], gif_path):
+    images[0].save(
+        gif_path,
+        save_all=True,
+        append_images=images[1:],
+        duration=200,  # Duration between frames in milliseconds
+        loop=0,
+    )  # Loop forever
+    return gif_path
+
+
+@spaces.GPU
+def generate(
+    source_image,
+    prompt,
+    negative_prompt,
+    seed,
+    num_inference_steps,
+    num_gol_steps,
+    gol_grid_dim,
+    img_size,
+    controlnet_conditioning_scale,
+    strength,
+    guidance_scale,
+):
+
+    # Compute the Game of Life first
+    gol = GameOfLife()
+    gol.set_random_state(dim=(gol_grid_dim, gol_grid_dim), p=0.5, seed=seed)
+    gol.generate_n_steps(n=num_gol_steps)
+
+    gol_grids = gol.game_history
+
+    # Generate the gif for the original Game of Life
+    gol_images = [
+        generate_image_from_grid(grid, img_size=img_size) for grid in gol_grids
+    ]
+    path_gol_gif = make_gif(gol_images, "gol_original.gif")
+
+    # Generate the gif for the ControlNet Game of Life
+    controlnet_images = generate_all_images(
+        gol_grids=gol_grids,
+        source_image=source_image,
+        num_inference_steps=num_inference_steps,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        strength=strength,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        seed=seed,
+        guidance_scale=guidance_scale,
+        img_size=img_size,
+    )
+
+    path_gol_controlnet = make_gif(controlnet_images, "gol_controlnet.gif")
+
+    return path_gol_controlnet, path_gol_gif
+
+
+source_image = gr.Image(label="Source Image", type="pil", value="sky-gol-image.jpeg")
+
+output_controlnet = gr.Image(label="ControlNet Game of Life")
+output_gol = gr.Image(label="Original Game of Life")
+prompt = gr.Textbox(
+    label="Prompt", value="clear sky with clouds, high quality, background 4k"
+)
+negative_prompt = gr.Textbox(
+    label="Negative Prompt",
+    value="ugly, disfigured, low quality, blurry, nsfw, qr code",
+)
+seed = gr.Number(label="Seed", value=42)
+num_inference_steps = gr.Number(label="Controlnet Inference Steps", value=50)
+num_gol_steps = gr.Slider(
+    label="Number of Game of Life Steps",
+    minimum=2,
+    maximum=100,
+    step=1,
+    value=40,
+)
+gol_grid_dim = gr.Number(
+    label="Game of Life Grid Dimension",
+    value=10,
+)
+
+img_size = gr.Number(label="Image Size (pixels)", value=512)
+controlnet_conditioning_scale = gr.Slider(
+    label="Controlnet Conditioning Scale", minimum=0.1, maximum=10.0, value=2.0
+)
+strength = gr.Slider(label="Strength", minimum=0.1, maximum=1.0, value=0.9)
+guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=100, value=20)
+
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        source_image,
+        prompt,
+        negative_prompt,
+        seed,
+        num_inference_steps,
+        num_gol_steps,
+        gol_grid_dim,
+        img_size,
+        controlnet_conditioning_scale,
+        strength,
+        guidance_scale,
+    ],
+    outputs=[output_controlnet, output_gol],
+)
+demo.launch()

controlnet.py

@@ -0,0 +1,67 @@
+import torch
+from diffusers import (
+    StableDiffusionControlNetImg2ImgPipeline,
+    ControlNetModel,
+    DDIMScheduler,
+)
+
+from PIL import Image
+
+
+class QRControlNet:
+
+    def __init__(self, device: str = "cuda"):
+
+        torch_dtype = torch.float16 if device == "cuda" else torch.float32
+        controlnet = ControlNetModel.from_pretrained(
+            "DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch_dtype
+        )
+
+        pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            controlnet=controlnet,
+            safety_checker=None,
+            torch_dtype=torch_dtype,
+        ).to(device)
+
+        if device == "cuda":
+            pipe.enable_xformers_memory_efficient_attention()
+            pipe.enable_model_cpu_offload()
+
+        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        self.pipe = pipe
+
+    def generate_image(
+        self,
+        source_image: Image,
+        control_image: Image,
+        prompt: str,
+        negative_prompt: str,
+        img_size=512,
+        num_inference_steps: int = 50,
+        guidance_scale: int = 20,
+        controlnet_conditioning_scale: float = 3.0,
+        strength=0.9,
+        seed=42,
+        **kwargs
+    ):
+
+        width = height = img_size
+        generator = torch.manual_seed(seed)
+
+        image = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=source_image,
+            control_image=control_image,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,  # 3.0,
+            generator=generator,
+            strength=strength,
+            num_inference_steps=num_inference_steps,
+            **kwargs
+        )
+
+        return image.images[0]

game_of_life.py

@@ -0,0 +1,97 @@
+import numpy as np
+from scipy import signal
+from PIL import Image
+
+from utils import generate_image_from_grid
+
+
+class GameOfLife:
+
+    def __init__(self):
+        """Initialize the game.
+        dim: dimensions of the board
+        p: probability of a cell being dead at init
+        seed: (optional) for reproducibility, set to None for a new random state
+        init_state: (optional) a np.array grid to start the game with
+        """
+        self.kernel = [
+            [1, 1, 1],
+            [1, 0, 1],
+            [1, 1, 1],
+        ]
+
+        self.kernel = np.ones((3, 3))
+        self.kernel[1, 1] = 0
+        self.game_history = []
+        self.state = None
+        self.step_counter = 0
+
+    def set_random_state(self, dim=(100, 100), p=0.5, seed=None):
+        if seed:
+            np.random.seed(seed)
+        self.state = (np.random.random(dim) < p).astype("int")
+        self.game_history.append(self.state.copy())
+
+    def set_empty_state(self, dim=(100, 100)):
+        self.state = np.zeros(dim)
+        self.game_history.append(self.state.copy())
+
+    def set_state_from_array(self, array):
+        self.state = array.copy()
+        self.game_history.append(self.state.copy())
+
+    def count_neighbors(self):
+        """
+        Count the number of live neighbors each cell in self.state has with convolutions.
+        """
+        self.neighbors = signal.convolve2d(
+            self.state, self.kernel, boundary="fill", fillvalue=0, mode="same"
+        ).astype("int")
+
+    def place_blob(self, blob, i, j):
+        """Place a blob at coordinates i,j
+        blob: ndarray of zeros and ones
+        i: int
+        j: int
+        """
+        try:
+            self.state[i : i + blob.shape[0], j : j + blob.shape[1]] = blob
+        except:
+            print("Check bounds of box vs size of game!")
+
+    def step(self):
+        """Update the game based on conway game of life rules"""
+        # Count the number of neighbors via convolution
+        self.count_neighbors()
+
+        # Copy of initial state
+        self.new_state = self.state
+
+        # Rebirth if cell is dead and has three live neighbors
+        self.new_state += np.logical_and(self.neighbors == 3, self.state == 0)
+
+        # Death if cell has less than 2 neighbors
+        self.new_state -= np.logical_and(self.neighbors < 2, self.state == 1)
+
+        # Death if cell has more than 3 neighbors
+        self.new_state -= np.logical_and(self.neighbors > 3, self.state == 1)
+
+        # Update game state
+        self.state = self.new_state
+
+        # Save state to history
+        self.game_history.append(self.state.copy())
+
+        # Update step counter
+        self.step_counter += 1
+
+    def generate_n_steps(self, n):
+        for _ in range(n):
+            self.step()
+
+            if np.array_equal(self.game_history[-1], self.game_history[-2]):
+                # If the game is stable, break
+                break
+
+    def generate_image(self, grid: np.array, img_size: int = 512) -> Image:
+        return generate_image_from_grid(grid, img_size)

requirements.txt

@@ -0,0 +1,7 @@
+torch 
+diffusers 
+pillow 
+transformers 
+accelerate 
+xformers 
+numpy

utils.py

@@ -0,0 +1,48 @@
+from PIL import Image
+import numpy as np
+import requests
+from io import BytesIO
+
+
+def resize_image(input_image: Image, resolution: int):
+    input_image = input_image.convert("RGB")
+    W, H = input_image.size
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    H = int(round(H / 64.0)) * 64
+    W = int(round(W / 64.0)) * 64
+    img = input_image.resize((W, H), resample=Image.LANCZOS)
+    return img
+
+
+def load_image(url):
+
+    # Mimic a browser request
+    headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
+    }
+
+    response = requests.get(url, headers=headers)
+    img = Image.open(BytesIO(response.content))
+    return img
+
+
+def generate_image_from_grid(grid: np.array, img_size: int = 512) -> Image:
+    """Generate an iamge from a grid of 0s and 1s"""
+    n = len(grid)
+    cell_pixel_size = img_size // n
+
+    # Create a new image with white background
+    img = Image.new("RGB", (img_size, img_size), "white")
+    pixels = img.load()
+
+    for i in range(n):
+        for j in range(n):
+            # Color a cell black if 0 or white if 1
+            color = (0, 0, 0) if grid[i][j] == 0 else (255, 255, 255)
+            for x in range(cell_pixel_size):
+                for y in range(cell_pixel_size):
+                    pixels[j * cell_pixel_size + x, i * cell_pixel_size + y] = color
+
+    return img