mit-han-lab / anycost-gan

[CVPR 2021] Anycost GANs for Interactive Image Synthesis and Editing
https://hanlab.mit.edu/projects/anycost-gan/
MIT License
776 stars 98 forks source link

生成图全是灰色 #18

Closed jxust01 closed 2 years ago

jxust01 commented 2 years ago
import torch
import numpy as np
import os, random
from PIL import Image
from tqdm import tqdm
from models.dynamic_channel import set_uniform_channel_ratio, reset_generator
import models
import time
import cv2
config = 'anycost-ffhq-config-f'
device = 'cuda:2'

class Face_Editor():
    def __init__(self):
        self.init_model()

    def init_model(self):
        self.anycost_channel = 1.0
        self.anycost_resolution = 1024
        self.generator = models.get_pretrained('generator', config).to(device)
        self.generator.eval()

    def sample(self):
        torch.manual_seed(1601)
        # latent = torch.randn(1, 1, 512, device=device)
        # mean_style = self.generator.mean_style(10000)
        # self.input_kwargs = {'styles':latent, 'return_rgbs':True, 'truncation':0.5,
        #                      'truncation_style':mean_style, 'randomize_noise':False}
        # style = torch.randn(1, 18, 512, device=device)
        style = np.load('/simple/zlp1/masters/anycost-gan/assets/demo_ori/projected_latents/00_ryan.npy')
        style = torch.from_numpy(style).view(1, -1, 512).to(device)
        self.input_kwargs = {'styles': style,
                            'noise': None, 'randomize_noise': False, 'input_is_style': True}
        image = self.generate_image()
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv2.imshow('image', image)
        cv2.waitKey(0)

    def generate_image(self):
        def image_to_np(x):
            assert x.shape[0] == 1
            x = x.squeeze(0).permute(1, 2, 0)
            x = (x + 1) * 0.5  # 0-1
            x = (x * 255).cpu().numpy().astype('uint8')
            return x

        with torch.no_grad():
            print(self.input_kwargs)
            out = self.generator(**self.input_kwargs)[0].clamp(-1, 1)
            out = image_to_np(out)
            return out

if __name__ == '__main__':
    FE = Face_Editor()
    FE.sample()
jxust01 commented 2 years ago

image

tonylins commented 2 years ago

Hi, please change device = 'cuda:2' to device = 'cuda', and it should work fine.

jxust01 commented 2 years ago

Hi, please change device = 'cuda:2' to device = 'cuda', and it should work fine.

thanks,it works,but why can't set the device = 'cuda:2'?

tonylins commented 2 years ago

Hi, please change device = 'cuda:2' to device = 'cuda', and it should work fine.

thanks,it works,but why can't set the device = 'cuda:2'?

I think it is a bug in the PyTorch custom CUDA kernel implementation. You can also use device='cuda2' by adding FORCE_NATIVE=1 to the environmental variables (see here https://github.com/mit-han-lab/anycost-gan/blob/1177beafaf5ed486156f78f411c34e62f0d46f14/models/ops.py#L8)

jxust01 commented 2 years ago

Thank you very much for your answers, looking forward to creating awesome things with anycost-gan