Delete src/image_generation.py

src/image_generation.py

@@ -1,366 +0,0 @@
-
-# Copyright (c) Facebook, Inc. and its affiliates.
-# Modified from: https://github.com/facebookresearch/detectron2/blob/master/demo/demo.py
-from transformers import pipeline
-import torchvision
-from PIL import Image
-from models.t2i_pipeline import StableDiffusionPipelineSpatialAware
-import torchvision.io as vision_io
-import torch.nn.functional as F
-import torch
-import tqdm
-import numpy as np
-import cv2
-import warnings
-import time
-import tempfile
-import argparse
-import glob
-import multiprocessing as mp
-import os
-import random
-
-# fmt: off
-import sys
-sys.path.insert(1, os.path.join(sys.path[0], '..'))
-# fmt: on
-
-
-warnings.filterwarnings("ignore")
-
-# constants
-WINDOW_NAME = "demo"
-
-
-def generate_image(pipe, overall_prompt, latents, get_latents=False, num_inference_steps=50, fg_masks=None,
-                   fg_masked_latents=None, frozen_steps=0, frozen_prompt=None, custom_attention_mask=None, fg_prompt=None):
-    '''
-    Main function that calls the image diffusion model
-    latent: input_noise from where it starts the generation
-    get_latents: if True, returns the latents for each frame
-    '''
-
-    image = pipe(overall_prompt, latents=latents, num_inference_steps=num_inference_steps, frozen_mask=fg_masks,
-                 frozen_steps=frozen_steps, latents_all_input=fg_masked_latents, frozen_prompt=frozen_prompt, custom_attention_mask=custom_attention_mask, output_type='pil', 
-                 fg_prompt=fg_prompt, make_attention_mask_2d=True, attention_mask_block_diagonal=True).images[0]
-    torch.save(image, "img.pt")
-
-    if get_latents:
-        video_latents = pipe(overall_prompt, latents=latents,
-                             num_inference_steps=num_inference_steps, output_type="latent").images
-        torch.save(video_latents, "img_latents.pt")
-        return image, video_latents
-
-    return image
-
-
-def save_frames(path):
-    video, audio, video_info = vision_io.read_video(
-        f"demo3/{path}.mp4", pts_unit='sec')
-
-    # Number of frames
-    num_frames = video.size(0)
-
-    # Save each frame
-    os.makedirs(f"demo3/{path}", exist_ok=True)
-    for i in range(num_frames):
-        frame = video[i, :, :, :].numpy()
-        # Convert from C x H x W to H x W x C and from torch tensor to PIL Image
-        # frame = frame.permute(1, 2, 0).numpy()
-        img = Image.fromarray(frame.astype('uint8'))
-        img.save(f"demo3/{path}/frame_{i:04d}.png")
-
-
-def create_boxes():
-    img_width = 96
-    img_height = 96
-
-    # initialize bboxes list
-    sbboxes = []
-
-    # object dimensions
-    for object_size in [20, 30, 40, 50, 60]:
-        obj_width, obj_height = object_size, object_size
-
-        # starting position
-        start_x = 3
-        start_y = 4
-
-        # calculate total size occupied by the objects in the grid
-        total_obj_width = 3 * obj_width
-        total_obj_height = 3 * obj_height
-
-        # determine horizontal and vertical spacings
-        spacing_horizontal = (img_width - total_obj_width - start_x) // 2
-        spacing_vertical = (img_height - total_obj_height - start_y) // 2
-
-        for i in range(3):
-            for j in range(3):
-                x_start = start_x + i * (obj_width + spacing_horizontal)
-                y_start = start_y + j * (obj_height + spacing_vertical)
-                # Corrected to img_width to include the last pixel
-                x_end = min(x_start + obj_width, img_width)
-                # Corrected to img_height to include the last pixel
-                y_end = min(y_start + obj_height, img_height)
-                sbboxes.append([x_start, y_start, x_end, y_end])
-
-    mask_id = 0
-    masks_list = []
-
-    for sbbox in sbboxes:
-        smask = torch.zeros(1, 1, 96, 96)
-        smask[0, 0, sbbox[1]:sbbox[3], sbbox[0]:sbbox[2]] = 1.0
-        masks_list.append(smask)
-        # torchvision.utils.save_image(smask, f"{SAVE_DIR}/masks/mask_{mask_id}.png")  # save masks as images
-        mask_id += 1
-
-    return masks_list
-
-
-def objects_list():
-    objects_settings = [
-        ("apple", "on a table"),
-        ("ball", "in a park"),
-        ("cat", "on a couch"),
-        ("dog", "in a backyard"),
-        ("elephant", "in a jungle"),
-        ("fountain pen", "on a desk"),
-        ("guitar", "on a stage"),
-        ("helicopter", "in the sky"),
-        ("island", "in the sea"),
-        ("jar", "on a shelf"),
-        ("kite", "in the sky"),
-        ("lamp", "in a room"),
-        ("motorbike", "on a road"),
-        ("notebook", "on a table"),
-        ("owl", "on a tree"),
-        ("piano", "in a hall"),
-        ("queen", "in a castle"),
-        ("robot", "in a lab"),
-        ("snake", "in a forest"),
-        ("tent", "in the mountains"),
-        ("umbrella", "on a beach"),
-        ("violin", "in an orchestra"),
-        ("wheel", "in a garage"),
-        ("xylophone", "in a music class"),
-        ("yacht", "in a marina"),
-        ("zebra", "in a savannah"),
-        ("aeroplane", "in the clouds"),
-        ("bridge", "over a river"),
-        ("computer", "in an office"),
-        ("dragon", "in a cave"),
-        ("egg", "in a nest"),
-        ("flower", "in a garden"),
-        ("globe", "in a library"),
-        ("hat", "on a rack"),
-        ("ice cube", "in a glass"),
-        ("jewelry", "in a box"),
-        ("kangaroo", "in a desert"),
-        ("lion", "in a den"),
-        ("mug", "on a counter"),
-        ("nest", "on a branch"),
-        ("octopus", "in the ocean"),
-        ("parrot", "in a rainforest"),
-        ("quilt", "on a bed"),
-        ("rose", "in a vase"),
-        ("ship", "in a dock"),
-        ("train", "on the tracks"),
-        ("utensils", "in a kitchen"),
-        ("vase", "on a window sill"),
-        ("watch", "in a store"),
-        ("x-ray", "in a hospital"),
-        ("yarn", "in a basket"),
-        ("zeppelin", "above a city"),
-    ]
-    objects_settings.extend([
-        ("muffin", "on a bakery shelf"),
-        ("notebook", "on a student's desk"),
-        ("owl", "in a tree"),
-        ("piano", "in a concert hall"),
-        ("quill", "on parchment"),
-        ("robot", "in a factory"),
-        ("snake", "in the grass"),
-        ("telescope", "in an observatory"),
-        ("umbrella", "at the beach"),
-        ("violin", "in an orchestra"),
-        ("whale", "in the ocean"),
-        ("xylophone", "in a music store"),
-        ("yacht", "in a marina"),
-        ("zebra", "on a savanna"),
-
-        # Kitchen items
-        ("spoon", "in a drawer"),
-        ("plate", "in a cupboard"),
-        ("cup", "on a shelf"),
-        ("frying pan", "on a stove"),
-        ("jar", "in the refrigerator"),
-
-        # Office items
-        ("computer", "in an office"),
-        ("printer", "by a desk"),
-        ("chair", "around a conference table"),
-        ("lamp", "on a workbench"),
-        ("calendar", "on a wall"),
-
-        # Outdoor items
-        ("bicycle", "on a street"),
-        ("tent", "in a campsite"),
-        ("fire", "in a fireplace"),
-        ("mountain", "in the distance"),
-        ("river", "through the woods"),
-
-
-        # and so on ...
-    ])
-
-    # To expedite the generation, you can combine themes and objects:
-
-    themes = [
-        ("wild animals", ["tiger", "lion", "cheetah",
-         "giraffe", "hippopotamus"], "in the wild"),
-        ("household items", ["sofa", "tv", "clock",
-         "vase", "photo frame"], "in a living room"),
-        ("clothes", ["shirt", "pants", "shoes",
-         "hat", "jacket"], "in a wardrobe"),
-        ("musical instruments", ["drum", "trumpet",
-         "harp", "saxophone", "tuba"], "in a band"),
-        ("cosmic entities", ["planet", "star",
-         "comet", "nebula", "asteroid"], "in space"),
-        # ... add more themes
-    ]
-
-    # Using the themes to extend our list
-    for theme_name, theme_objects, theme_location in themes:
-        for theme_object in theme_objects:
-            objects_settings.append((theme_object, theme_location))
-
-    # Sports equipment
-    objects_settings.extend([
-        ("basketball", "on a court"),
-        ("golf ball", "on a golf course"),
-        ("tennis racket", "on a tennis court"),
-        ("baseball bat", "in a stadium"),
-        ("hockey stick", "on an ice rink"),
-        ("football", "on a field"),
-        ("skateboard", "in a skatepark"),
-        ("boxing gloves", "in a boxing ring"),
-        ("ski", "on a snowy slope"),
-        ("surfboard", "on a beach shore"),
-    ])
-
-    # Toys and games
-    objects_settings.extend([
-        ("teddy bear", "on a child's bed"),
-        ("doll", "in a toy store"),
-        ("toy car", "on a carpet"),
-        ("board game", "on a table"),
-        ("yo-yo", "in a child's hand"),
-        ("kite", "in the sky on a windy day"),
-        ("Lego bricks", "on a construction table"),
-        ("jigsaw puzzle", "partially completed"),
-        ("rubik's cube", "on a shelf"),
-        ("action figure", "on display"),
-    ])
-
-    # Transportation
-    objects_settings.extend([
-        ("bus", "at a bus stop"),
-        ("motorcycle", "on a road"),
-        ("helicopter", "landing on a pad"),
-        ("scooter", "on a sidewalk"),
-        ("train", "at a station"),
-        ("bicycle", "parked by a post"),
-        ("boat", "in a harbor"),
-        ("tractor", "on a farm"),
-        ("airplane", "taking off from a runway"),
-        ("submarine", "below sea level"),
-    ])
-
-    # Medieval theme
-    objects_settings.extend([
-        ("castle", "on a hilltop"),
-        ("knight", "riding a horse"),
-        ("bow and arrow", "in an archery range"),
-        ("crown", "in a treasure chest"),
-        ("dragon", "flying over mountains"),
-        ("shield", "next to a warrior"),
-        ("dagger", "on a wooden table"),
-        ("torch", "lighting a dark corridor"),
-        ("scroll", "sealed with wax"),
-        ("cauldron", "with bubbling potion"),
-    ])
-
-    # Modern technology
-    objects_settings.extend([
-        ("smartphone", "on a charger"),
-        ("laptop", "in a cafe"),
-        ("headphones", "around a neck"),
-        ("camera", "on a tripod"),
-        ("drone", "flying over a park"),
-        ("USB stick", "plugged into a computer"),
-        ("watch", "on a wrist"),
-        ("microphone", "on a podcast desk"),
-        ("tablet", "with a digital pen"),
-        ("VR headset", "ready for gaming"),
-    ])
-
-    # Nature
-    objects_settings.extend([
-        ("tree", "in a forest"),
-        ("flower", "in a garden"),
-        ("mountain", "on a horizon"),
-        ("cloud", "in a blue sky"),
-        ("waterfall", "in a scenic location"),
-        ("beach", "next to an ocean"),
-        ("cactus", "in a desert"),
-        ("volcano", "erupting with lava"),
-        ("coral", "under the sea"),
-        ("moon", "in a night sky"),
-    ])
-
-    prompts = [f"A {obj} {setting}" for obj, setting in objects_settings]
-
-    return objects_settings
-
-
-if __name__ == "__main__":
-    SAVE_DIR = "/scr/image/"
-    save_path = "img43-att_mask"
-    torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    random_latents = torch.randn(
-        [1, 4, 96, 96], generator=torch.Generator().manual_seed(1)).to(torch_device)
-
-    try:
-        pipe = StableDiffusionPipelineSpatialAware.from_pretrained(
-            "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float, variant="fp32", cache_dir="/gscratch/scrubbed/anasery/").to(torch_device)
-    except:
-        pipe = StableDiffusionPipelineSpatialAware.from_pretrained(
-            "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float, variant="fp32").to(torch_device)
-
-    fg_object = "apple"  # fg object stores the object to be masked
-    # overall prompt stores the prompt
-    overall_prompt = f"An {fg_object} on plate"
-    os.makedirs(f"{SAVE_DIR}/{overall_prompt}", exist_ok=True)
-
-    masks_list = create_boxes()
-
-    # torch.save(f"{overall_prompt}+masked", "prompt.pt")
-    obj_settings = objects_list()  # 166
-    for obj_setting in obj_settings[120:]:
-        fg_object = obj_setting[0]
-        overall_prompt = f"A {obj_setting[0]} {obj_setting[1]}"
-        print(overall_prompt)
-
-        # randomly select 10 numbers from range len of masks_list
-        selected_mask_ids = random.sample(range(len(masks_list)), 3)
-        for mask_id in selected_mask_ids:
-            os.makedirs(
-                f"{SAVE_DIR}/{overall_prompt}/mask{mask_id}", exist_ok=True)
-            torchvision.utils.save_image(
-                masks_list[mask_id][0][0], f"{SAVE_DIR}/{overall_prompt}/mask{mask_id}/mask.png")
-            for frozen_steps in range(0, 5):
-                img = generate_image(pipe, overall_prompt, random_latents, get_latents=False, num_inference_steps=50, fg_masks=masks_list[mask_id].to(
-                    torch_device), fg_masked_latents=None, frozen_steps=frozen_steps, frozen_prompt=None, fg_prompt=fg_object)
-
-                img.save(
-                    f"{SAVE_DIR}/{overall_prompt}/mask{mask_id}/{frozen_steps}.png")