Delete test_seesr_sam.py

test_seesr_sam.py

@@ -1,328 +0,0 @@
-'''
- * SeeSR: Towards Semantics-Aware Real-World Image Super-Resolution 
- * Modified from diffusers by Rongyuan Wu
- * 24/12/2023
-'''
-import os
-import sys
-sys.path.append(os.getcwd())
-import cv2
-import glob
-import argparse
-import numpy as np
-from PIL import Image
-
-import torch
-import torch.utils.checkpoint
-
-from accelerate import Accelerator
-from accelerate.logging import get_logger
-from accelerate.utils import set_seed
-from diffusers import AutoencoderKL, DDPMScheduler
-from diffusers.utils import check_min_version
-from diffusers.utils.import_utils import is_xformers_available
-from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor
-
-from pipelines.pipeline_seesr import StableDiffusionControlNetPipeline
-from utils.misc import load_dreambooth_lora
-from utils.wavelet_color_fix import wavelet_color_fix, adain_color_fix
-
-from ram.models.ram_lora import ram
-from ram import inference_ram as inference
-from ram import get_transform
-
-from typing import Mapping, Any
-from torchvision import transforms
-import torch.nn as nn
-import torch.nn.functional as F
-from torchvision import transforms
-
-sys.path.insert(0, '/media/ssd8T/wyw/code/SeeSR/sam2')
-from sam2.build_sam import build_sam2
-from sam2.sam2_image_predictor import SAM2ImagePredictor
-
-logger = get_logger(__name__, log_level="INFO")
-
-
-tensor_transforms = transforms.Compose([
-                transforms.ToTensor(),
-            ])
-
-ram_transforms = transforms.Compose([
-            transforms.Resize((384, 384)),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
-
-sam_mean = [0.485, 0.456, 0.406]
-sam_std = [0.229, 0.224, 0.225]
-sam_transforms = transforms.Compose([
-    transforms.Resize((1024, 1024)),
-    transforms.Normalize(mean=sam_mean, std=sam_std)
-])
-        
-        
-def load_state_dict_diffbirSwinIR(model: nn.Module, state_dict: Mapping[str, Any], strict: bool=False) -> None:
-    state_dict = state_dict.get("state_dict", state_dict)
-    
-    is_model_key_starts_with_module = list(model.state_dict().keys())[0].startswith("module.")
-    is_state_dict_key_starts_with_module = list(state_dict.keys())[0].startswith("module.")
-    
-    if (
-        is_model_key_starts_with_module and
-        (not is_state_dict_key_starts_with_module)
-    ):
-        state_dict = {f"module.{key}": value for key, value in state_dict.items()}
-    if (
-        (not is_model_key_starts_with_module) and
-        is_state_dict_key_starts_with_module
-    ):
-        state_dict = {key[len("module."):]: value for key, value in state_dict.items()}
-    
-    model.load_state_dict(state_dict, strict=strict)
-
-
-def load_seesr_pipeline(args, accelerator, enable_xformers_memory_efficient_attention):
-    
-    from models.controlnet import ControlNetModel
-    from models.unet_2d_condition import UNet2DConditionModel
-
-    # Load scheduler, tokenizer and models.
-    
-    scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_path, subfolder="scheduler")
-    text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_path, subfolder="text_encoder")
-    tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_path, subfolder="tokenizer")
-    vae = AutoencoderKL.from_pretrained(args.pretrained_model_path, subfolder="vae")
-    feature_extractor = CLIPImageProcessor.from_pretrained(f"{args.pretrained_model_path}/feature_extractor")
-    unet = UNet2DConditionModel.from_pretrained(args.seesr_model_path, subfolder="unet")
-    controlnet = ControlNetModel.from_pretrained(args.seesr_model_path, subfolder="controlnet")
-    
-    # Freeze vae and text_encoder
-    vae.requires_grad_(False)
-    text_encoder.requires_grad_(False)
-    unet.requires_grad_(False)
-    controlnet.requires_grad_(False)
-
-    if enable_xformers_memory_efficient_attention:
-        if is_xformers_available():
-            unet.enable_xformers_memory_efficient_attention()
-            controlnet.enable_xformers_memory_efficient_attention()
-        else:
-            raise ValueError("xformers is not available. Make sure it is installed correctly")
-
-    # Get the validation pipeline
-    validation_pipeline = StableDiffusionControlNetPipeline(
-        vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, feature_extractor=feature_extractor, 
-        unet=unet, controlnet=controlnet, scheduler=scheduler, safety_checker=None, requires_safety_checker=False,
-    )
-    # def count_parameters(model):
-    #     return sum(p.numel() for p in model.parameters())
-
-    # # 计算各子模块的参数量
-    # unet_params = count_parameters(unet)
-    # controlnet_params = count_parameters(controlnet)
-    # vae_params = count_parameters(vae)
-    # text_encoder_params = count_parameters(text_encoder)
-
-    # # 总参数量
-    # total_params = unet_params + controlnet_params + vae_params + text_encoder_params
-    # print(f"UNet 参数量: {unet_params/ 1e9:.2f}B")
-    # print(f"ControlNet 参数量: {controlnet_params/ 1e9:.2f}B")
-    # print(f"VAE 参数��: {vae_params/ 1e9:.2f}B")
-    # print(f"Text Encoder 参数量: {text_encoder_params/ 1e9:.2f}B")
-    # print(f"总参数量: {total_params/ 1e9:.2f}B")
-    # exit()
-    
-    
-    validation_pipeline._init_tiled_vae(encoder_tile_size=args.vae_encoder_tiled_size, decoder_tile_size=args.vae_decoder_tiled_size)
-
-    # For mixed precision training we cast the text_encoder and vae weights to half-precision
-    # as these models are only used for inference, keeping weights in full precision is not required.
-    weight_dtype = torch.float32
-    if accelerator.mixed_precision == "fp16":
-        weight_dtype = torch.float16
-    elif accelerator.mixed_precision == "bf16":
-        weight_dtype = torch.bfloat16
-
-    # Move text_encode and vae to gpu and cast to weight_dtype
-    text_encoder.to(accelerator.device, dtype=weight_dtype)
-    vae.to(accelerator.device, dtype=weight_dtype)
-    unet.to(accelerator.device, dtype=weight_dtype)
-    controlnet.to(accelerator.device, dtype=weight_dtype)
-
-    return validation_pipeline
-
-def load_tag_model(args, device='cuda'):#  0.47B
-    
-    model = ram(pretrained='/media/ssd8T/ly/SeeSR/preset/models/ram_swin_large_14m.pth',
-                pretrained_condition=args.ram_ft_path,
-                image_size=384,
-                vit='swin_l')
-    model.eval()
-    model.to(device)
-    # def count_parameters(model):
-    #     return sum(p.numel() for p in model.parameters())
-    # total_params = count_parameters(model)
-    # print(f"总参数量: {total_params/ 1e9:.2f}B")
-    # exit()
-    return model
-
-def get_validation_prompt(args, image, model, device='cuda'):
-    validation_prompt = ""
- 
-    lq = tensor_transforms(image).unsqueeze(0).to(device)
-    lq = ram_transforms(lq)
-    res = inference(lq, model)
-    ram_encoder_hidden_states = model.generate_image_embeds(lq)
-
-    validation_prompt = f"{res[0]}, {args.prompt},"
-
-    return validation_prompt, ram_encoder_hidden_states
-
-def load_sam_model(device='cuda'):#0.03B
-    
-    sam2_checkpoint = "./preset/models/sam2.1_hiera_tiny.pt"
-    model_cfg = "configs/sam2.1/sam2.1_hiera_t.yaml"
-    sam2 = build_sam2(model_cfg, sam2_checkpoint, mode="eval", device='cuda', apply_postprocessing=False)
-    sam2.to(device)
-    SAM = SAM2ImagePredictor(sam2)
-
-    return SAM
-
-def get_sam_embedding(image, model, device='cuda'):
- 
-    lq = tensor_transforms(image).unsqueeze(0).to(device)
-    lq = sam_transforms(lq)
-    sam_encoder_hidden_states = model.generate_image_embedding(lq)
-
-    return sam_encoder_hidden_states
-
-def main(args, enable_xformers_memory_efficient_attention=True,):
-    txt_path = os.path.join(args.output_dir, 'txt')
-    os.makedirs(txt_path, exist_ok=True)
-
-    accelerator = Accelerator(
-        mixed_precision=args.mixed_precision,
-    )
-
-    # If passed along, set the training seed now.
-    if args.seed is not None:
-        set_seed(args.seed)
-
-    # Handle the output folder creation
-    if accelerator.is_main_process:
-        os.makedirs(args.output_dir, exist_ok=True)
-
-    # We need to initialize the trackers we use, and also store our configuration.
-    # The trackers initializes automatically on the main process.
-    if accelerator.is_main_process:
-        accelerator.init_trackers("SeeSR")
-
-    pipeline = load_seesr_pipeline(args, accelerator, enable_xformers_memory_efficient_attention)
-    model = load_tag_model(args, accelerator.device)
-    sam_model = load_sam_model(accelerator.device)
- 
-    if accelerator.is_main_process:
-        generator = torch.Generator(device=accelerator.device)
-        if args.seed is not None:
-            generator.manual_seed(args.seed)
-
-        if os.path.isdir(args.image_path):
-            image_names = sorted(glob.glob(f'{args.image_path}/*.*'))
-        else:
-            image_names = [args.image_path]
-        # image_names = image_names[:10]
-
-        for image_idx, image_name in enumerate(image_names[:]):
-            print(f'================== process {image_idx} imgs... ===================')
-            validation_image = Image.open(image_name).convert("RGB")
-
-            validation_prompt, _ = get_validation_prompt(args, validation_image, model)
-            validation_prompt += args.added_prompt # clean, extremely detailed, best quality, sharp, clean
-            negative_prompt = args.negative_prompt #dirty, messy, low quality, frames, deformed, 
-            
-            sam_encoder_hidden_states = get_sam_embedding(validation_image,sam_model)
-            
-            if args.save_prompts:
-                txt_save_path = f"{txt_path}/{os.path.basename(image_name).split('.')[0]}.txt"
-                file = open(txt_save_path, "w")
-                file.write(validation_prompt)
-                file.close()
-            print(f'{validation_prompt}')
-
-            ori_width, ori_height = validation_image.size
-            resize_flag = False
-            rscale = args.upscale
-            if ori_width < args.process_size//rscale or ori_height < args.process_size//rscale:
-                scale = (args.process_size//rscale)/min(ori_width, ori_height)
-                tmp_image = validation_image.resize((int(scale*ori_width), int(scale*ori_height)))
-
-                validation_image = tmp_image
-                resize_flag = True
-
-            validation_image = validation_image.resize((validation_image.size[0]*rscale, validation_image.size[1]*rscale))
-            validation_image = validation_image.resize((validation_image.size[0]//8*8, validation_image.size[1]//8*8))
-            width, height = validation_image.size
-            resize_flag = True #
-
-            print(f'input size: {height}x{width}')
-
-            for sample_idx in range(args.sample_times):
-                os.makedirs(f'{args.output_dir}/sample{str(sample_idx).zfill(2)}/', exist_ok=True)
-
-            for sample_idx in range(args.sample_times):  
-                with torch.autocast("cuda"):
-                    image = pipeline(
-                            validation_prompt, validation_image, num_inference_steps=args.num_inference_steps, generator=generator, height=height, width=width,
-                            guidance_scale=args.guidance_scale, negative_prompt=negative_prompt, conditioning_scale=args.conditioning_scale,
-                            start_point=args.start_point, ram_encoder_hidden_states=sam_encoder_hidden_states,
-                            latent_tiled_size=args.latent_tiled_size, latent_tiled_overlap=args.latent_tiled_overlap,
-                            args=args,
-                        ).images[0]
-                
-                if args.align_method == 'nofix':
-                    image = image
-                else:
-                    if args.align_method == 'wavelet':
-                        image = wavelet_color_fix(image, validation_image)
-                    elif args.align_method == 'adain':
-                        image = adain_color_fix(image, validation_image)
-
-                if resize_flag: 
-                    image = image.resize((ori_width*rscale, ori_height*rscale))
-                    
-                name, ext = os.path.splitext(os.path.basename(image_name))
-                
-                image.save(f'{args.output_dir}/sample{str(sample_idx).zfill(2)}/{name}.png')
-    
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--seesr_model_path", type=str, default=None)
-    parser.add_argument("--ram_ft_path", type=str, default=None)
-    parser.add_argument("--pretrained_model_path", type=str, default=None)
-    parser.add_argument("--prompt", type=str, default="") # user can add self-prompt to improve the results
-    parser.add_argument("--added_prompt", type=str, default="clean, high-resolution, 8k")
-    parser.add_argument("--negative_prompt", type=str, default="dotted, noise, blur, lowres, smooth")
-    parser.add_argument("--image_path", type=str, default=None)
-    parser.add_argument("--output_dir", type=str, default=None)
-    parser.add_argument("--mixed_precision", type=str, default="fp16") # no/fp16/bf16
-    parser.add_argument("--guidance_scale", type=float, default=5.5)
-    parser.add_argument("--conditioning_scale", type=float, default=1.0)
-    parser.add_argument("--blending_alpha", type=float, default=1.0)
-    parser.add_argument("--num_inference_steps", type=int, default=50)
-    parser.add_argument("--process_size", type=int, default=512)
-    parser.add_argument("--vae_decoder_tiled_size", type=int, default=224) # latent size, for 24G
-    parser.add_argument("--vae_encoder_tiled_size", type=int, default=1024) # image size, for 13G
-    parser.add_argument("--latent_tiled_size", type=int, default=96) 
-    parser.add_argument("--latent_tiled_overlap", type=int, default=32) 
-    parser.add_argument("--upscale", type=int, default=4)
-    parser.add_argument("--seed", type=int, default=None)
-    parser.add_argument("--sample_times", type=int, default=1)
-    parser.add_argument("--align_method", type=str, choices=['wavelet', 'adain', 'nofix'], default='adain')
-    parser.add_argument("--start_steps", type=int, default=999) # defaults set to 999.
-    parser.add_argument("--start_point", type=str, choices=['lr', 'noise'], default='lr') # LR Embedding Strategy, choose 'lr latent + 999 steps noise' as diffusion start point. 
-    parser.add_argument("--save_prompts", action='store_true')
-    args = parser.parse_args()
-    main(args)
-
-
-