Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from diffusers import PriorTransformer, UNet2DConditionModel, KandinskyV22Pipeline | |
| from huggingface_hub import hf_hub_download | |
| from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor, CLIPTokenizer, CLIPTextModelWithProjection | |
| from model import pops_utils | |
| from model.pipeline_pops import pOpsPipeline | |
| kandinsky_prior_repo: str = 'kandinsky-community/kandinsky-2-2-prior' | |
| kandinsky_decoder_repo: str = 'kandinsky-community/kandinsky-2-2-decoder' | |
| prior_texture_repo: str = 'models/texturing/learned_prior.pth' | |
| prior_instruct_repo: str = 'models/instruct/learned_prior.pth' | |
| prior_scene_repo: str = 'models/scene/learned_prior.pth' | |
| prior_repo = "pOpsPaper/operators" | |
| # gpu = torch.device('cuda') | |
| # cpu = torch.device('cpu') | |
| class PopsPipelines: | |
| def __init__(self): | |
| weight_dtype = torch.float16 | |
| self.weight_dtype = weight_dtype | |
| device = 'cpu' #torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.device = 'cuda' #device | |
| self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(kandinsky_prior_repo, | |
| subfolder='image_encoder', | |
| torch_dtype=weight_dtype).eval() | |
| self.image_encoder.requires_grad_(False) | |
| self.image_processor = CLIPImageProcessor.from_pretrained(kandinsky_prior_repo, | |
| subfolder='image_processor') | |
| self.tokenizer = CLIPTokenizer.from_pretrained(kandinsky_prior_repo, subfolder='tokenizer') | |
| self.text_encoder = CLIPTextModelWithProjection.from_pretrained(kandinsky_prior_repo, | |
| subfolder='text_encoder', | |
| torch_dtype=weight_dtype).eval().to(device) | |
| # Load full model for vis | |
| self.unet = UNet2DConditionModel.from_pretrained(kandinsky_decoder_repo, | |
| subfolder='unet').to(torch.float16).to(device) | |
| self.decoder = KandinskyV22Pipeline.from_pretrained(kandinsky_decoder_repo, unet=self.unet, | |
| torch_dtype=torch.float16) | |
| self.decoder = self.decoder.to(device) | |
| self.priors_dict = { | |
| 'texturing':{'repo':prior_texture_repo}, | |
| 'instruct': {'repo': prior_instruct_repo}, | |
| 'scene': {'repo':prior_scene_repo} | |
| } | |
| for prior_type in self.priors_dict: | |
| prior_path = self.priors_dict[prior_type]['repo'] | |
| prior = PriorTransformer.from_pretrained( | |
| kandinsky_prior_repo, subfolder="prior" | |
| ) | |
| # Load from huggingface | |
| prior_path = hf_hub_download(repo_id=prior_repo, filename=str(prior_path)) | |
| prior_state_dict = torch.load(prior_path, map_location=device) | |
| prior.load_state_dict(prior_state_dict, strict=False) | |
| prior.eval() | |
| prior = prior.to(weight_dtype) | |
| prior_pipeline = pOpsPipeline.from_pretrained(kandinsky_prior_repo, | |
| prior=prior, | |
| image_encoder=self.image_encoder, | |
| torch_dtype=torch.float16) | |
| self.priors_dict[prior_type]['pipeline'] = prior_pipeline | |
| def process_image(self, input_path): | |
| if input_path is None: | |
| return None | |
| image_pil = Image.open(input_path).convert("RGB").resize((512, 512)) | |
| image = torch.Tensor(self.image_processor(image_pil)['pixel_values'][0]).to(self.device).unsqueeze(0).to( | |
| self.weight_dtype) | |
| return image | |
| def process_text(self, text): | |
| self.text_encoder.to('cuda') | |
| text_inputs = self.tokenizer( | |
| text, | |
| padding="max_length", | |
| max_length=self.tokenizer.model_max_length, | |
| truncation=True, | |
| return_tensors="pt", | |
| ) | |
| mask = text_inputs.attention_mask.bool() # [0] | |
| text_encoder_output = self.text_encoder(text_inputs.input_ids.to(self.device)) | |
| text_encoder_hidden_states = text_encoder_output.last_hidden_state | |
| text_encoder_concat = text_encoder_hidden_states[:, :mask.sum().item()] | |
| self.text_encoder.to('cpu') | |
| return text_encoder_concat | |
| def run_binary(self, input_a, input_b, prior_type): | |
| # Move pipeline to GPU | |
| pipeline = self.priors_dict[prior_type]['pipeline'] | |
| pipeline.to('cuda') | |
| self.image_encoder.to('cuda') | |
| input_image_embeds, input_hidden_state = pops_utils.preprocess(input_a, input_b, | |
| self.image_encoder, | |
| pipeline.prior.clip_mean.detach(), | |
| pipeline.prior.clip_std.detach()) | |
| negative_input_embeds = torch.zeros_like(input_image_embeds) | |
| negative_hidden_states = torch.zeros_like(input_hidden_state) | |
| guidance_scale = 1.0 | |
| if prior_type == 'texturing': | |
| guidance_scale = 8.0 | |
| img_emb = pipeline(input_embeds=input_image_embeds, input_hidden_states=input_hidden_state, | |
| negative_input_embeds=negative_input_embeds, | |
| negative_input_hidden_states=negative_hidden_states, | |
| num_inference_steps=25, | |
| num_images_per_prompt=1, | |
| guidance_scale=guidance_scale) | |
| # Optional | |
| if prior_type == 'scene': | |
| # Scene is the closet to what avg represents for a background image so incorporate that as well | |
| mean_emb = 0.5 * input_hidden_state[:, 0] + 0.5 * input_hidden_state[:, 1] | |
| mean_emb = (mean_emb * pipeline.prior.clip_std) + pipeline.prior.clip_mean | |
| alpha = 0.4 | |
| img_emb.image_embeds = (1 - alpha) * img_emb.image_embeds + alpha * mean_emb | |
| # Move pipeline to CPU | |
| pipeline.to('cpu') | |
| self.image_encoder.to('cpu') | |
| return img_emb | |
| def run_instruct(self, input_a, text): | |
| text_encodings = self.process_text(text) | |
| # Move pipeline to GPU | |
| instruct_pipeline = self.priors_dict['instruct']['pipeline'] | |
| instruct_pipeline.to('cuda') | |
| self.image_encoder.to('cuda') | |
| input_image_embeds, input_hidden_state = pops_utils.preprocess(input_a, None, | |
| self.image_encoder, | |
| instruct_pipeline.prior.clip_mean.detach(), instruct_pipeline.prior.clip_std.detach(), | |
| concat_hidden_states=text_encodings) | |
| negative_input_embeds = torch.zeros_like(input_image_embeds) | |
| negative_hidden_states = torch.zeros_like(input_hidden_state) | |
| img_emb = instruct_pipeline(input_embeds=input_image_embeds, input_hidden_states=input_hidden_state, | |
| negative_input_embeds=negative_input_embeds, | |
| negative_input_hidden_states=negative_hidden_states, | |
| num_inference_steps=25, | |
| num_images_per_prompt=1, | |
| guidance_scale=1.0) | |
| # Move pipeline to CPU | |
| instruct_pipeline.to('cpu') | |
| self.image_encoder.to('cpu') | |
| return img_emb | |
| def render(self, img_emb): | |
| self.decoder.to('cuda') | |
| images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=img_emb.negative_image_embeds, | |
| num_inference_steps=50, height=512, | |
| width=512, guidance_scale=4).images | |
| self.decoder.to('cpu') | |
| return images[0] | |
| def run_instruct_texture(self, image_object_path, text_instruct, image_texture_path): | |
| # Process both inputs | |
| image_object = self.process_image(image_object_path) | |
| image_texture = self.process_image(image_texture_path) | |
| if image_object is None: | |
| raise gr.Error('Object image is required') | |
| current_emb = None | |
| if image_texture is None: | |
| instruct_input = image_object | |
| else: | |
| # Run texturing | |
| current_emb = self.run_binary(input_a=image_object, input_b=image_texture,prior_type='texturing') | |
| instruct_input = current_emb.image_embeds | |
| if text_instruct != '': | |
| current_emb = self.run_instruct(input_a=instruct_input, text=text_instruct) | |
| if current_emb is None: | |
| raise gr.Error('At least one of the inputs is required') | |
| # Render as image | |
| image = self.render(current_emb) | |
| return image | |
| def run_texture_scene(self, image_object_path, image_texture_path, image_scene_path): | |
| # Process both inputs | |
| image_object = self.process_image(image_object_path) | |
| image_texture = self.process_image(image_texture_path) | |
| image_scene = self.process_image(image_scene_path) | |
| if image_object is None: | |
| raise gr.Error('Object image is required') | |
| current_emb = None | |
| if image_texture is None: | |
| scene_input = image_object | |
| else: | |
| # Run texturing | |
| current_emb = self.run_binary(input_a=image_object, input_b=image_scene,prior_type='scene') | |
| scene_input = current_emb.image_embeds | |
| # Run scene | |
| if image_scene is not None: | |
| current_emb = self.run_binary(input_a=scene_input, input_b=image_texture,prior_type='texturing') | |
| if current_emb is None: | |
| raise gr.Error('At least one of the images is required') | |
| # Render as image | |
| image = self.render(current_emb) | |
| return image | |