rijulpaul/ddpm-ffhq-128

Diffusion Model for Human Face Generation (128x128)

Useful Links:

Overview

This project uses a DDPM based Diffusion techinique to learn to generate human faces from noise. The dataset used is flickr-faces-ffhq

Running the Model

# pipeline.py
from diffusers import DiffusionPipeline
import torch
from PIL import Image

class Pipeline:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        model_id = "rijulpaul/ddpm-ffhq-128"

        self.pipe = DiffusionPipeline.from_pretrained(
            model_id,
        ).to(self.device)

    def generate(self, seed=None, num_inference_steps=20):
        generator = None

        if seed is not None and seed != -1:
            generator = torch.Generator(device=self.device).manual_seed(int(seed))

        image = self.pipe(
            num_inference_steps=num_inference_steps,
            generator=generator
        ).images[0]

        return image

pipe = Pipeline()

def generate_image(seed, steps):
    if seed == -1:
        seed = random.randint(0, 1_000_000)

    img = pipe.generate(seed=seed, num_inference_steps=steps)

    if not isinstance(img, Image.Image):
      img = Image.fromarray(img)

    path = "output.png"
    img.save(path)

seed = 182341
steps = 1000
generate_image(seed,steps)