diff --git a/garment_adapter/attention_processor.py b/garment_adapter/attention_processor.py index 1d0f397..80dbc30 100644 --- a/garment_adapter/attention_processor.py +++ b/garment_adapter/attention_processor.py @@ -17,16 +17,16 @@ def __init__(self): super().__init__() def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - scale: float = 1.0, - attn_store=None, - do_classifier_free_guidance=None, - enable_cloth_guidance=None + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + attn_store=None, + do_classifier_free_guidance=None, + enable_cloth_guidance=None, ) -> torch.Tensor: residual = hidden_states @@ -39,22 +39,32 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape + ) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) query = attn.to_q(hidden_states, *args) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states, *args) value = attn.to_v(encoder_hidden_states, *args) @@ -73,7 +83,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -90,16 +102,16 @@ def __init__(self, name, type="read"): self.type = type def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - scale: float = 1.0, - attn_store=None, - do_classifier_free_guidance=None, - enable_cloth_guidance=None + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + attn_store=None, + do_classifier_free_guidance=None, + enable_cloth_guidance=None, ) -> torch.Tensor: if self.type == "read": attn_store[self.name] = hidden_states @@ -108,7 +120,9 @@ def __call__( if do_classifier_free_guidance: empty_copy = torch.zeros_like(ref_hidden_states) if enable_cloth_guidance: - ref_hidden_states = torch.cat([empty_copy, ref_hidden_states, ref_hidden_states]) + ref_hidden_states = torch.cat( + [empty_copy, ref_hidden_states, ref_hidden_states] + ) else: ref_hidden_states = torch.cat([empty_copy, ref_hidden_states]) hidden_states = torch.cat([hidden_states, ref_hidden_states], dim=1) @@ -125,22 +139,32 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape + ) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) query = attn.to_q(hidden_states, *args) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states, *args) value = attn.to_v(encoder_hidden_states, *args) @@ -162,7 +186,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -180,19 +206,21 @@ class AttnProcessor2_0(nn.Module): def __init__(self): super().__init__() if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - scale: float = 1.0, - attn_store=None, - do_classifier_free_guidance=None, - enable_cloth_guidance=None + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + attn_store=None, + do_classifier_free_guidance=None, + enable_cloth_guidance=None, ) -> torch.FloatTensor: residual = hidden_states if attn.spatial_norm is not None: @@ -202,20 +230,30 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape ) if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size + ) # scaled_dot_product_attention expects attention_mask shape to be # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.view( + batch_size, attn.heads, -1, attention_mask.shape[-1] + ) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) args = () if USE_PEFT_BACKEND else (scale,) query = attn.to_q(hidden_states, *args) @@ -223,7 +261,9 @@ def __call__( if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states, *args) value = attn.to_v(encoder_hidden_states, *args) @@ -242,7 +282,9 @@ def __call__( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) hidden_states = hidden_states.to(query.dtype) # linear proj @@ -251,7 +293,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -265,21 +309,23 @@ class REFAttnProcessor2_0(nn.Module): def __init__(self, name, type="read"): super().__init__() if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) self.name = name self.type = type def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - scale: float = 1.0, - attn_store=None, - do_classifier_free_guidance=False, - enable_cloth_guidance=True + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + attn_store=None, + do_classifier_free_guidance=False, + enable_cloth_guidance=True, ) -> torch.FloatTensor: if self.type == "read": attn_store[self.name] = hidden_states @@ -288,7 +334,9 @@ def __call__( if do_classifier_free_guidance: empty_copy = torch.zeros_like(ref_hidden_states) if enable_cloth_guidance: - ref_hidden_states = torch.cat([empty_copy, ref_hidden_states, ref_hidden_states]) + ref_hidden_states = torch.cat( + [empty_copy, ref_hidden_states, ref_hidden_states] + ) else: ref_hidden_states = torch.cat([empty_copy, ref_hidden_states]) hidden_states = torch.cat([hidden_states, ref_hidden_states], dim=1) @@ -302,20 +350,30 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape ) if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size + ) # scaled_dot_product_attention expects attention_mask shape to be # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.view( + batch_size, attn.heads, -1, attention_mask.shape[-1] + ) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) args = () if USE_PEFT_BACKEND else (scale,) query = attn.to_q(hidden_states, *args) @@ -323,7 +381,9 @@ def __call__( if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states, *args) value = attn.to_v(encoder_hidden_states, *args) @@ -342,7 +402,9 @@ def __call__( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) hidden_states = hidden_states.to(query.dtype) if self.type == "write": @@ -353,7 +415,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -366,20 +430,22 @@ class REFAnimateDiffAttnProcessor2_0(nn.Module): def __init__(self, name, type="read"): super().__init__() if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) self.name = name self.type = type def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - scale: float = 1.0, - attn_store=None, - do_classifier_free_guidance=False, + self, + attn: Attention, + hidden_states: torch.FloatTensor, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + attn_store=None, + do_classifier_free_guidance=False, ) -> torch.FloatTensor: if self.type == "read": attn_store[self.name] = hidden_states @@ -387,11 +453,21 @@ def __call__( ref_hidden_states = attn_store[self.name] if do_classifier_free_guidance: empty_copy = torch.zeros_like(ref_hidden_states) - ref_hidden_states = torch.cat([empty_copy, ref_hidden_states, ref_hidden_states]) + ref_hidden_states = torch.cat( + [empty_copy, ref_hidden_states, ref_hidden_states] + ) if hidden_states.shape[0] % ref_hidden_states.shape[0] != 0: raise ValueError("not evenly divisible") # ref_hidden_states = ref_hidden_states*1.05 - hidden_states = torch.cat([hidden_states, ref_hidden_states.repeat(hidden_states.shape[0] // ref_hidden_states.shape[0], 1, 1)], dim=1) + hidden_states = torch.cat( + [ + hidden_states, + ref_hidden_states.repeat( + hidden_states.shape[0] // ref_hidden_states.shape[0], 1, 1 + ), + ], + dim=1, + ) else: raise ValueError("unsupport type") residual = hidden_states @@ -402,20 +478,30 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape ) if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size + ) # scaled_dot_product_attention expects attention_mask shape to be # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.view( + batch_size, attn.heads, -1, attention_mask.shape[-1] + ) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) args = () if USE_PEFT_BACKEND else (scale,) query = attn.to_q(hidden_states, *args) @@ -423,7 +509,9 @@ def __call__( if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states, *args) value = attn.to_v(encoder_hidden_states, *args) @@ -442,7 +530,9 @@ def __call__( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) hidden_states = hidden_states.to(query.dtype) if self.type == "write": @@ -453,7 +543,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -463,7 +555,6 @@ def __call__( class IPAttnProcessor(nn.Module): - def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4): super().__init__() @@ -472,19 +563,23 @@ def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens= self.scale = scale self.num_tokens = num_tokens - self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) + self.to_k_ip = nn.Linear( + cross_attention_dim or hidden_size, hidden_size, bias=False + ) + self.to_v_ip = nn.Linear( + cross_attention_dim or hidden_size, hidden_size, bias=False + ) def __call__( - self, - attn, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - temb=None, - attn_store=None, - do_classifier_free_guidance=None, - enable_cloth_guidance=None + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + temb=None, + attn_store=None, + do_classifier_free_guidance=None, + enable_cloth_guidance=None, ): residual = hidden_states @@ -495,15 +590,23 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape + ) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size ) - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) query = attn.to_q(hidden_states) @@ -517,7 +620,9 @@ def __call__( encoder_hidden_states[:, end_pos:, :], ) if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) @@ -550,7 +655,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual @@ -561,31 +668,36 @@ def __call__( class IPAttnProcessor2_0(torch.nn.Module): - def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, num_tokens=4): super().__init__() if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") + raise ImportError( + "AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0." + ) self.hidden_size = hidden_size self.cross_attention_dim = cross_attention_dim self.scale = scale self.num_tokens = num_tokens - self.to_k_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) - self.to_v_ip = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) + self.to_k_ip = nn.Linear( + cross_attention_dim or hidden_size, hidden_size, bias=False + ) + self.to_v_ip = nn.Linear( + cross_attention_dim or hidden_size, hidden_size, bias=False + ) def __call__( - self, - attn, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - temb=None, - attn_store=None, - do_classifier_free_guidance=None, - enable_cloth_guidance=None + self, + attn, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + temb=None, + attn_store=None, + do_classifier_free_guidance=None, + enable_cloth_guidance=None, ): residual = hidden_states @@ -596,20 +708,30 @@ def __call__( if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + hidden_states = hidden_states.view( + batch_size, channel, height * width + ).transpose(1, 2) batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + hidden_states.shape + if encoder_hidden_states is None + else encoder_hidden_states.shape ) if attention_mask is not None: - attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + attention_mask = attn.prepare_attention_mask( + attention_mask, sequence_length, batch_size + ) # scaled_dot_product_attention expects attention_mask shape to be # (batch, heads, source_length, target_length) - attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) + attention_mask = attention_mask.view( + batch_size, attn.heads, -1, attention_mask.shape[-1] + ) if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose( + 1, 2 + ) query = attn.to_q(hidden_states) @@ -623,7 +745,9 @@ def __call__( encoder_hidden_states[:, end_pos:, :], ) if attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + encoder_hidden_states = attn.norm_encoder_hidden_states( + encoder_hidden_states + ) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) @@ -642,7 +766,9 @@ def __call__( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + hidden_states = hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) hidden_states = hidden_states.to(query.dtype) # for ip-adapter @@ -661,7 +787,9 @@ def __call__( self.attn_map = query @ ip_key.transpose(-2, -1).softmax(dim=-1) # print(self.attn_map.shape) - ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) + ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape( + batch_size, -1, attn.heads * head_dim + ) ip_hidden_states = ip_hidden_states.to(query.dtype) hidden_states = hidden_states + self.scale * ip_hidden_states @@ -672,7 +800,9 @@ def __call__( hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + hidden_states = hidden_states.transpose(-1, -2).reshape( + batch_size, channel, height, width + ) if attn.residual_connection: hidden_states = hidden_states + residual diff --git a/garment_adapter/garment_diffusion.py b/garment_adapter/garment_diffusion.py index 4fb7ff3..68d3c14 100644 --- a/garment_adapter/garment_diffusion.py +++ b/garment_adapter/garment_diffusion.py @@ -8,13 +8,17 @@ if is_torch2_available(): from .attention_processor import REFAttnProcessor2_0 as REFAttnProcessor from .attention_processor import AttnProcessor2_0 as AttnProcessor - from .attention_processor import REFAnimateDiffAttnProcessor2_0 as REFAnimateDiffAttnProcessor + from .attention_processor import ( + REFAnimateDiffAttnProcessor2_0 as REFAnimateDiffAttnProcessor, + ) else: from .attention_processor import REFAttnProcessor, AttnProcessor class ClothAdapter: - def __init__(self, sd_pipe, ref_path, device, enable_cloth_guidance, set_seg_model=True): + def __init__( + self, sd_pipe, ref_path, device, enable_cloth_guidance, set_seg_model=True + ): self.enable_cloth_guidance = enable_cloth_guidance self.device = device self.pipe = sd_pipe.to(self.device) @@ -22,7 +26,13 @@ def __init__(self, sd_pipe, ref_path, device, enable_cloth_guidance, set_seg_mod ref_unet = copy.deepcopy(sd_pipe.unet) if ref_unet.config.in_channels == 9: - ref_unet.conv_in = torch.nn.Conv2d(4, 320, ref_unet.conv_in.kernel_size, ref_unet.conv_in.stride, ref_unet.conv_in.padding) + ref_unet.conv_in = torch.nn.Conv2d( + 4, + 320, + ref_unet.conv_in.kernel_size, + ref_unet.conv_in.stride, + ref_unet.conv_in.padding, + ) ref_unet.register_to_config(in_channels=4) state_dict = {} with safe_open(ref_path, framework="pt", device="cpu") as f: @@ -36,8 +46,10 @@ def __init__(self, sd_pipe, ref_path, device, enable_cloth_guidance, set_seg_mod self.set_seg_model() self.attn_store = {} - def set_seg_model(self, ): - checkpoint_path = 'checkpoints/cloth_segm.pth' + def set_seg_model( + self, + ): + checkpoint_path = "checkpoints/cloth_segm.pth" self.seg_net = load_seg_model(checkpoint_path, device=self.device) def set_adapter(self, unet, type): @@ -50,23 +62,25 @@ def set_adapter(self, unet, type): unet.set_attn_processor(attn_procs) def generate( - self, - cloth_image, - cloth_mask_image=None, - prompt=None, - a_prompt="best quality, high quality", - num_images_per_prompt=4, - negative_prompt=None, - seed=-1, - guidance_scale=7.5, - cloth_guidance_scale=2.5, - num_inference_steps=20, - height=512, - width=384, - **kwargs, + self, + cloth_image, + cloth_mask_image=None, + prompt=None, + a_prompt="best quality, high quality", + num_images_per_prompt=4, + negative_prompt=None, + seed=-1, + guidance_scale=7.5, + cloth_guidance_scale=2.5, + num_inference_steps=20, + height=512, + width=384, + **kwargs, ): if cloth_mask_image is None: - cloth_mask_image = generate_mask(cloth_image, net=self.seg_net, device=self.device) + cloth_mask_image = generate_mask( + cloth_image, net=self.seg_net, device=self.device + ) cloth = prepare_image(cloth_image, height, width) cloth_mask = prepare_mask(cloth_mask_image, height, width) @@ -76,7 +90,9 @@ def generate( prompt = "a photography of a model" prompt = prompt + ", " + a_prompt if negative_prompt is None: - negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" + negative_prompt = ( + "monochrome, lowres, bad anatomy, worst quality, low quality" + ) with torch.inference_mode(): prompt_embeds, negative_prompt_embeds = self.pipe.encode_prompt( @@ -86,11 +102,26 @@ def generate( do_classifier_free_guidance=True, negative_prompt=negative_prompt, ) - prompt_embeds_null = self.pipe.encode_prompt([""], device=self.device, num_images_per_prompt=num_images_per_prompt, do_classifier_free_guidance=False)[0] - cloth_embeds = self.pipe.vae.encode(cloth).latent_dist.mode() * self.pipe.vae.config.scaling_factor - self.ref_unet(torch.cat([cloth_embeds] * num_images_per_prompt), 0, prompt_embeds_null, cross_attention_kwargs={"attn_store": self.attn_store}) + prompt_embeds_null = self.pipe.encode_prompt( + [""], + device=self.device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=False, + )[0] + cloth_embeds = ( + self.pipe.vae.encode(cloth).latent_dist.mode() + * self.pipe.vae.config.scaling_factor + ) + self.ref_unet( + torch.cat([cloth_embeds] * num_images_per_prompt), + 0, + prompt_embeds_null, + cross_attention_kwargs={"attn_store": self.attn_store}, + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) if self.enable_cloth_guidance: images = self.pipe( prompt_embeds=prompt_embeds, @@ -101,7 +132,11 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0, "enable_cloth_guidance": self.enable_cloth_guidance}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + "enable_cloth_guidance": self.enable_cloth_guidance, + }, **kwargs, ).images else: @@ -113,37 +148,58 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0, "enable_cloth_guidance": self.enable_cloth_guidance}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + "enable_cloth_guidance": self.enable_cloth_guidance, + }, **kwargs, ).images return images, cloth_mask_image def generate_inpainting( - self, - cloth_image, - cloth_mask_image=None, - num_images_per_prompt=4, - seed=-1, - cloth_guidance_scale=2.5, - num_inference_steps=20, - height=512, - width=384, - **kwargs, + self, + cloth_image, + cloth_mask_image=None, + num_images_per_prompt=4, + seed=-1, + cloth_guidance_scale=2.5, + num_inference_steps=20, + height=512, + width=384, + **kwargs, ): if cloth_mask_image is None: - cloth_mask_image = generate_mask(cloth_image, net=self.seg_net, device=self.device) + cloth_mask_image = generate_mask( + cloth_image, net=self.seg_net, device=self.device + ) cloth = prepare_image(cloth_image, height, width) cloth_mask = prepare_mask(cloth_mask_image, height, width) cloth = (cloth * cloth_mask).to(self.device, dtype=torch.float16) with torch.inference_mode(): - prompt_embeds_null = self.pipe.encode_prompt([""], device=self.device, num_images_per_prompt=num_images_per_prompt, do_classifier_free_guidance=False)[0] - cloth_embeds = self.pipe.vae.encode(cloth).latent_dist.mode() * self.pipe.vae.config.scaling_factor - self.ref_unet(torch.cat([cloth_embeds] * num_images_per_prompt), 0, prompt_embeds_null, cross_attention_kwargs={"attn_store": self.attn_store}) + prompt_embeds_null = self.pipe.encode_prompt( + [""], + device=self.device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=False, + )[0] + cloth_embeds = ( + self.pipe.vae.encode(cloth).latent_dist.mode() + * self.pipe.vae.config.scaling_factor + ) + self.ref_unet( + torch.cat([cloth_embeds] * num_images_per_prompt), + 0, + prompt_embeds_null, + cross_attention_kwargs={"attn_store": self.attn_store}, + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) images = self.pipe( prompt_embeds=prompt_embeds_null, cloth_guidance_scale=cloth_guidance_scale, @@ -151,7 +207,11 @@ def generate_inpainting( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": cloth_guidance_scale > 1.0, "enable_cloth_guidance": False}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": cloth_guidance_scale > 1.0, + "enable_cloth_guidance": False, + }, **kwargs, ).images @@ -164,7 +224,9 @@ def __init__(self, sd_pipe, pipe_path, ref_path, device, set_seg_model=True): self.pipe = sd_pipe.to(self.device) self.set_adapter(self.pipe.unet, "write") - ref_unet = UNet2DConditionModel.from_pretrained(pipe_path, subfolder='unet', torch_dtype=sd_pipe.dtype) + ref_unet = UNet2DConditionModel.from_pretrained( + pipe_path, subfolder="unet", torch_dtype=sd_pipe.dtype + ) state_dict = {} with safe_open(ref_path, framework="pt", device="cpu") as f: for key in f.keys(): @@ -177,8 +239,10 @@ def __init__(self, sd_pipe, pipe_path, ref_path, device, set_seg_model=True): self.set_seg_model() self.attn_store = {} - def set_seg_model(self, ): - checkpoint_path = 'checkpoints/cloth_segm.pth' + def set_seg_model( + self, + ): + checkpoint_path = "checkpoints/cloth_segm.pth" self.seg_net = load_seg_model(checkpoint_path, device=self.device) def set_adapter(self, unet, type): @@ -191,23 +255,25 @@ def set_adapter(self, unet, type): unet.set_attn_processor(attn_procs) def generate( - self, - cloth_image, - cloth_mask_image=None, - prompt=None, - a_prompt="best quality, high quality", - num_images_per_prompt=4, - negative_prompt=None, - seed=-1, - guidance_scale=7.5, - cloth_guidance_scale=3., - num_inference_steps=20, - height=512, - width=384, - **kwargs, + self, + cloth_image, + cloth_mask_image=None, + prompt=None, + a_prompt="best quality, high quality", + num_images_per_prompt=4, + negative_prompt=None, + seed=-1, + guidance_scale=7.5, + cloth_guidance_scale=3.0, + num_inference_steps=20, + height=512, + width=384, + **kwargs, ): if cloth_mask_image is None: - cloth_mask_image = generate_mask(cloth_image, net=self.seg_net, device=self.device) + cloth_mask_image = generate_mask( + cloth_image, net=self.seg_net, device=self.device + ) cloth = prepare_image(cloth_image, height, width) cloth_mask = prepare_mask(cloth_mask_image, height, width) @@ -227,11 +293,26 @@ def generate( do_classifier_free_guidance=True, negative_prompt=negative_prompt, ) - prompt_embeds_null = self.pipe.encode_prompt([""], device=self.device, num_images_per_prompt=num_images_per_prompt, do_classifier_free_guidance=False)[0] - cloth_embeds = self.pipe.vae.encode(cloth).latent_dist.mode() * self.pipe.vae.config.scaling_factor - self.ref_unet(torch.cat([cloth_embeds] * num_images_per_prompt), 0, prompt_embeds_null, cross_attention_kwargs={"attn_store": self.attn_store}) + prompt_embeds_null = self.pipe.encode_prompt( + [""], + device=self.device, + num_images_per_prompt=num_images_per_prompt, + do_classifier_free_guidance=False, + )[0] + cloth_embeds = ( + self.pipe.vae.encode(cloth).latent_dist.mode() + * self.pipe.vae.config.scaling_factor + ) + self.ref_unet( + torch.cat([cloth_embeds] * num_images_per_prompt), + 0, + prompt_embeds_null, + cross_attention_kwargs={"attn_store": self.attn_store}, + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) frames = self.pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -241,7 +322,10 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + }, **kwargs, ).frames diff --git a/garment_adapter/garment_ipadapter_faceid.py b/garment_adapter/garment_ipadapter_faceid.py index b54c2fd..e99d69e 100644 --- a/garment_adapter/garment_ipadapter_faceid.py +++ b/garment_adapter/garment_ipadapter_faceid.py @@ -26,15 +26,15 @@ class FacePerceiverResampler(torch.nn.Module): def __init__( - self, - *, - dim=768, - depth=4, - dim_head=64, - heads=16, - embedding_dim=1280, - output_dim=768, - ff_mult=4, + self, + *, + dim=768, + depth=4, + dim_head=64, + heads=16, + embedding_dim=1280, + output_dim=768, + ff_mult=4, ): super().__init__() @@ -83,7 +83,13 @@ def forward(self, id_embeds): class ProjPlusModel(torch.nn.Module): - def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, clip_embeddings_dim=1280, num_tokens=4): + def __init__( + self, + cross_attention_dim=768, + id_embeddings_dim=512, + clip_embeddings_dim=1280, + num_tokens=4, + ): super().__init__() self.cross_attention_dim = cross_attention_dim @@ -117,7 +123,18 @@ def forward(self, id_embeds, clip_embeds, shortcut=False, scale=1.0): class IPAdapterFaceID: - def __init__(self, sd_pipe, ref_path, ip_ckpt, device, enable_cloth_guidance, num_tokens=4, n_cond=1, torch_dtype=torch.float16, set_seg_model=True): + def __init__( + self, + sd_pipe, + ref_path, + ip_ckpt, + device, + enable_cloth_guidance, + num_tokens=4, + n_cond=1, + torch_dtype=torch.float16, + set_seg_model=True, + ): self.enable_cloth_guidance = enable_cloth_guidance self.device = device self.ip_ckpt = ip_ckpt @@ -149,11 +166,16 @@ def __init__(self, sd_pipe, ref_path, ip_ckpt, device, enable_cloth_guidance, nu self.attn_store = {} def set_insightface(self): - self.app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) + self.app = FaceAnalysis( + name="buffalo_l", + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) self.app.prepare(ctx_id=0, det_size=(640, 640)) - def set_seg_model(self, ): - checkpoint_path = 'checkpoints/cloth_segm.pth' + def set_seg_model( + self, + ): + checkpoint_path = "checkpoints/cloth_segm.pth" self.seg_net = load_seg_model(checkpoint_path, device=self.device) def init_proj(self): @@ -177,7 +199,11 @@ def set_ip_adapter(self): unet = self.pipe.unet attn_procs = {} for name in unet.attn_processors.keys(): - cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim + cross_attention_dim = ( + None + if name.endswith("attn1.processor") + else unet.config.cross_attention_dim + ) if name.startswith("mid_block"): hidden_size = unet.config.block_out_channels[-1] elif name.startswith("up_blocks"): @@ -190,7 +216,10 @@ def set_ip_adapter(self): attn_procs[name] = REFAttnProcessor(name=name, type="write") else: attn_procs[name] = IPAttnProcessor( - hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0, num_tokens=self.num_tokens * self.n_cond, + hidden_size=hidden_size, + cross_attention_dim=cross_attention_dim, + scale=1.0, + num_tokens=self.num_tokens * self.n_cond, ).to(self.device, dtype=self.torch_dtype) unet.set_attn_processor(attn_procs) @@ -200,9 +229,13 @@ def load_ip_adapter(self): with safe_open(self.ip_ckpt, framework="pt", device="cpu") as f: for key in f.keys(): if key.startswith("image_proj."): - state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key) + state_dict["image_proj"][key.replace("image_proj.", "")] = ( + f.get_tensor(key) + ) elif key.startswith("ip_adapter."): - state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key) + state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = ( + f.get_tensor(key) + ) else: state_dict = torch.load(self.ip_ckpt, map_location="cpu") self.image_proj_model.load_state_dict(state_dict["image_proj"]) @@ -211,7 +244,6 @@ def load_ip_adapter(self): @torch.inference_mode() def get_image_embeds(self, faceid_embeds): - multi_face = False if faceid_embeds.dim() == 3: multi_face = True @@ -220,7 +252,9 @@ def get_image_embeds(self, faceid_embeds): faceid_embeds = faceid_embeds.to(self.device, dtype=self.torch_dtype) image_prompt_embeds = self.image_proj_model(faceid_embeds) - uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(faceid_embeds)) + uncond_image_prompt_embeds = self.image_proj_model( + torch.zeros_like(faceid_embeds) + ) if multi_face: c = image_prompt_embeds.size(-1) image_prompt_embeds = image_prompt_embeds.reshape(b, -1, c) @@ -234,22 +268,22 @@ def set_scale(self, scale): attn_processor.scale = scale def generate( - self, - cloth_image, - face_image, - cloth_mask=None, - prompt=None, - a_prompt="best quality, high quality", - negative_prompt=None, - num_samples=4, - seed=None, - guidance_scale=3., - cloth_guidance_scale=3., - num_inference_steps=30, - height=512, - width=384, - scale=1.0, - **kwargs, + self, + cloth_image, + face_image, + cloth_mask=None, + prompt=None, + a_prompt="best quality, high quality", + negative_prompt=None, + num_samples=4, + seed=None, + guidance_scale=3.0, + cloth_guidance_scale=3.0, + num_inference_steps=30, + height=512, + width=384, + scale=1.0, + **kwargs, ): faces = self.app.get(cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR)) try: @@ -258,7 +292,9 @@ def generate( return None if cloth_mask is None: - cloth_mask_image = generate_mask(cloth_image, net=self.seg_net, device=self.device) + cloth_mask_image = generate_mask( + cloth_image, net=self.seg_net, device=self.device + ) cloth = prepare_image(cloth_image, height, width) cloth_mask = prepare_mask(cloth_mask_image, height, width) @@ -272,20 +308,30 @@ def generate( prompt = "a photography of a model" prompt = prompt + ", " + a_prompt if negative_prompt is None: - negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" + negative_prompt = ( + "monochrome, lowres, bad anatomy, worst quality, low quality" + ) if not isinstance(prompt, List): prompt = [prompt] * num_prompts if not isinstance(negative_prompt, List): negative_prompt = [negative_prompt] * num_prompts - image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(faceid_embeds) + image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds( + faceid_embeds + ) bs_embed, seq_len, _ = image_prompt_embeds.shape image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1) - image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) + image_prompt_embeds = image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat( + 1, num_samples, 1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) with torch.inference_mode(): prompt_embeds_, negative_prompt_embeds_ = self.pipe.encode_prompt( @@ -296,13 +342,30 @@ def generate( negative_prompt=negative_prompt, ) prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1) - negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1) + negative_prompt_embeds = torch.cat( + [negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1 + ) - prompt_embeds_null = self.pipe.encode_prompt([""], device=self.device, num_images_per_prompt=num_samples, do_classifier_free_guidance=False)[0] - cloth_embeds = self.pipe.vae.encode(cloth).latent_dist.mode() * self.pipe.vae.config.scaling_factor - self.ref_unet(torch.cat([cloth_embeds] * num_samples), 0, prompt_embeds_null, cross_attention_kwargs={"attn_store": self.attn_store}) + prompt_embeds_null = self.pipe.encode_prompt( + [""], + device=self.device, + num_images_per_prompt=num_samples, + do_classifier_free_guidance=False, + )[0] + cloth_embeds = ( + self.pipe.vae.encode(cloth).latent_dist.mode() + * self.pipe.vae.config.scaling_factor + ) + self.ref_unet( + torch.cat([cloth_embeds] * num_samples), + 0, + prompt_embeds_null, + cross_attention_kwargs={"attn_store": self.attn_store}, + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) images = self.pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -311,7 +374,11 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0, "enable_cloth_guidance": self.enable_cloth_guidance}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + "enable_cloth_guidance": self.enable_cloth_guidance, + }, **kwargs, ).images @@ -319,7 +386,18 @@ def generate( class IPAdapterFaceIDPlus: - def __init__(self, sd_pipe, ref_path, image_encoder_path, ip_ckpt, device, enable_cloth_guidance, num_tokens=4, torch_dtype=torch.float16, set_seg_model=True): + def __init__( + self, + sd_pipe, + ref_path, + image_encoder_path, + ip_ckpt, + device, + enable_cloth_guidance, + num_tokens=4, + torch_dtype=torch.float16, + set_seg_model=True, + ): self.enable_cloth_guidance = enable_cloth_guidance self.device = device self.image_encoder_path = image_encoder_path @@ -331,9 +409,9 @@ def __init__(self, sd_pipe, ref_path, image_encoder_path, ip_ckpt, device, enabl self.set_ip_adapter() # load image encoder - self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.image_encoder_path).to( - self.device, dtype=self.torch_dtype - ) + self.image_encoder = CLIPVisionModelWithProjection.from_pretrained( + self.image_encoder_path + ).to(self.device, dtype=self.torch_dtype) self.clip_image_processor = CLIPImageProcessor() # image proj model self.image_proj_model = self.init_proj() @@ -356,11 +434,16 @@ def __init__(self, sd_pipe, ref_path, image_encoder_path, ip_ckpt, device, enabl self.attn_store = {} def set_insightface(self): - self.app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) + self.app = FaceAnalysis( + name="buffalo_l", + providers=["CUDAExecutionProvider", "CPUExecutionProvider"], + ) self.app.prepare(ctx_id=0, det_size=(640, 640)) - def set_seg_model(self, ): - checkpoint_path = 'checkpoints/cloth_segm.pth' + def set_seg_model( + self, + ): + checkpoint_path = "checkpoints/cloth_segm.pth" self.seg_net = load_seg_model(checkpoint_path, device=self.device) def init_proj(self): @@ -385,7 +468,11 @@ def set_ip_adapter(self): unet = self.pipe.unet attn_procs = {} for name in unet.attn_processors.keys(): - cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim + cross_attention_dim = ( + None + if name.endswith("attn1.processor") + else unet.config.cross_attention_dim + ) if name.startswith("mid_block"): hidden_size = unet.config.block_out_channels[-1] elif name.startswith("up_blocks"): @@ -398,7 +485,10 @@ def set_ip_adapter(self): attn_procs[name] = REFAttnProcessor(name=name, type="write") else: attn_procs[name] = IPAttnProcessor( - hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, scale=1.0, num_tokens=self.num_tokens, + hidden_size=hidden_size, + cross_attention_dim=cross_attention_dim, + scale=1.0, + num_tokens=self.num_tokens, ).to(self.device, dtype=self.torch_dtype) unet.set_attn_processor(attn_procs) @@ -408,9 +498,13 @@ def load_ip_adapter(self): with safe_open(self.ip_ckpt, framework="pt", device="cpu") as f: for key in f.keys(): if key.startswith("image_proj."): - state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key) + state_dict["image_proj"][key.replace("image_proj.", "")] = ( + f.get_tensor(key) + ) elif key.startswith("ip_adapter."): - state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key) + state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = ( + f.get_tensor(key) + ) else: state_dict = torch.load(self.ip_ckpt, map_location="cpu") self.image_proj_model.load_state_dict(state_dict["image_proj"]) @@ -419,16 +513,27 @@ def load_ip_adapter(self): @torch.inference_mode() def get_image_embeds(self, faceid_embeds, face_image, s_scale, shortcut): - clip_image = self.clip_image_processor(images=face_image, return_tensors="pt").pixel_values + clip_image = self.clip_image_processor( + images=face_image, return_tensors="pt" + ).pixel_values clip_image = clip_image.to(self.device, dtype=self.torch_dtype) - clip_image_embeds = self.image_encoder(clip_image, output_hidden_states=True).hidden_states[-2] + clip_image_embeds = self.image_encoder( + clip_image, output_hidden_states=True + ).hidden_states[-2] uncond_clip_image_embeds = self.image_encoder( torch.zeros_like(clip_image), output_hidden_states=True ).hidden_states[-2] faceid_embeds = faceid_embeds.to(self.device, dtype=self.torch_dtype) - image_prompt_embeds = self.image_proj_model(faceid_embeds, clip_image_embeds, shortcut=shortcut, scale=s_scale) - uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(faceid_embeds), uncond_clip_image_embeds, shortcut=shortcut, scale=s_scale) + image_prompt_embeds = self.image_proj_model( + faceid_embeds, clip_image_embeds, shortcut=shortcut, scale=s_scale + ) + uncond_image_prompt_embeds = self.image_proj_model( + torch.zeros_like(faceid_embeds), + uncond_clip_image_embeds, + shortcut=shortcut, + scale=s_scale, + ) return image_prompt_embeds, uncond_image_prompt_embeds def set_scale(self, scale): @@ -437,35 +542,39 @@ def set_scale(self, scale): attn_processor.scale = scale def generate( - self, - cloth_image, - face_image, - cloth_mask=None, - prompt=None, - a_prompt="best quality, high quality", - negative_prompt=None, - num_samples=4, - seed=None, - guidance_scale=2.5, - cloth_guidance_scale=2.5, - num_inference_steps=20, - height=512, - width=384, - scale=1.0, - s_scale=1., - shortcut=False, - **kwargs, + self, + cloth_image, + face_image, + cloth_mask=None, + prompt=None, + a_prompt="best quality, high quality", + negative_prompt=None, + num_samples=4, + seed=None, + guidance_scale=2.5, + cloth_guidance_scale=2.5, + num_inference_steps=20, + height=512, + width=384, + scale=1.0, + s_scale=1.0, + shortcut=False, + **kwargs, ): face_image = cv2.cvtColor(np.array(face_image), cv2.COLOR_RGB2BGR) faces = self.app.get(face_image) try: faceid_embeds = torch.from_numpy(faces[0].normed_embedding).unsqueeze(0) - face_image = face_align.norm_crop(face_image, landmark=faces[0].kps, image_size=224) + face_image = face_align.norm_crop( + face_image, landmark=faces[0].kps, image_size=224 + ) except: return None if cloth_mask is None: - cloth_mask_image = generate_mask(cloth_image, net=self.seg_net, device=self.device) + cloth_mask_image = generate_mask( + cloth_image, net=self.seg_net, device=self.device + ) cloth = prepare_image(cloth_image, height, width) cloth_mask = prepare_mask(cloth_mask_image, height, width) @@ -478,20 +587,30 @@ def generate( prompt = "a photography of a model" prompt = prompt + ", " + a_prompt if negative_prompt is None: - negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" + negative_prompt = ( + "monochrome, lowres, bad anatomy, worst quality, low quality" + ) if not isinstance(prompt, List): prompt = [prompt] * num_prompts if not isinstance(negative_prompt, List): negative_prompt = [negative_prompt] * num_prompts - image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(faceid_embeds, face_image, s_scale, shortcut) + image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds( + faceid_embeds, face_image, s_scale, shortcut + ) bs_embed, seq_len, _ = image_prompt_embeds.shape image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1) - image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) + image_prompt_embeds = image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat( + 1, num_samples, 1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) with torch.inference_mode(): prompt_embeds_, negative_prompt_embeds_ = self.pipe.encode_prompt( @@ -502,13 +621,30 @@ def generate( negative_prompt=negative_prompt, ) prompt_embeds = torch.cat([prompt_embeds_, image_prompt_embeds], dim=1) - negative_prompt_embeds = torch.cat([negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1) + negative_prompt_embeds = torch.cat( + [negative_prompt_embeds_, uncond_image_prompt_embeds], dim=1 + ) - prompt_embeds_null = self.pipe.encode_prompt([""], device=self.device, num_images_per_prompt=num_samples, do_classifier_free_guidance=False)[0] - cloth_embeds = self.pipe.vae.encode(cloth).latent_dist.mode() * self.pipe.vae.config.scaling_factor - self.ref_unet(torch.cat([cloth_embeds] * num_samples), 0, prompt_embeds_null, cross_attention_kwargs={"attn_store": self.attn_store}) + prompt_embeds_null = self.pipe.encode_prompt( + [""], + device=self.device, + num_images_per_prompt=num_samples, + do_classifier_free_guidance=False, + )[0] + cloth_embeds = ( + self.pipe.vae.encode(cloth).latent_dist.mode() + * self.pipe.vae.config.scaling_factor + ) + self.ref_unet( + torch.cat([cloth_embeds] * num_samples), + 0, + prompt_embeds_null, + cross_attention_kwargs={"attn_store": self.attn_store}, + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) if self.enable_cloth_guidance: images = self.pipe( prompt_embeds=prompt_embeds, @@ -519,7 +655,11 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0, "enable_cloth_guidance": self.enable_cloth_guidance}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + "enable_cloth_guidance": self.enable_cloth_guidance, + }, **kwargs, ).images else: @@ -531,7 +671,11 @@ def generate( generator=generator, height=height, width=width, - cross_attention_kwargs={"attn_store": self.attn_store, "do_classifier_free_guidance": guidance_scale > 1.0, "enable_cloth_guidance": self.enable_cloth_guidance}, + cross_attention_kwargs={ + "attn_store": self.attn_store, + "do_classifier_free_guidance": guidance_scale > 1.0, + "enable_cloth_guidance": self.enable_cloth_guidance, + }, **kwargs, ).images @@ -542,15 +686,15 @@ class IPAdapterFaceIDXL(IPAdapterFaceID): """SDXL""" def generate( - self, - faceid_embeds=None, - prompt=None, - negative_prompt=None, - scale=1.0, - num_samples=4, - seed=None, - num_inference_steps=30, - **kwargs, + self, + faceid_embeds=None, + prompt=None, + negative_prompt=None, + scale=1.0, + num_samples=4, + seed=None, + num_inference_steps=30, + **kwargs, ): self.set_scale(scale) @@ -559,20 +703,30 @@ def generate( if prompt is None: prompt = "best quality, high quality" if negative_prompt is None: - negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" + negative_prompt = ( + "monochrome, lowres, bad anatomy, worst quality, low quality" + ) if not isinstance(prompt, List): prompt = [prompt] * num_prompts if not isinstance(negative_prompt, List): negative_prompt = [negative_prompt] * num_prompts - image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(faceid_embeds) + image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds( + faceid_embeds + ) bs_embed, seq_len, _ = image_prompt_embeds.shape image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1) - image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) + image_prompt_embeds = image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat( + 1, num_samples, 1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) with torch.inference_mode(): ( @@ -587,9 +741,13 @@ def generate( negative_prompt=negative_prompt, ) prompt_embeds = torch.cat([prompt_embeds, image_prompt_embeds], dim=1) - negative_prompt_embeds = torch.cat([negative_prompt_embeds, uncond_image_prompt_embeds], dim=1) + negative_prompt_embeds = torch.cat( + [negative_prompt_embeds, uncond_image_prompt_embeds], dim=1 + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) images = self.pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -607,19 +765,19 @@ class IPAdapterFaceIDPlusXL(IPAdapterFaceIDPlus): """SDXL""" def generate( - self, - face_image=None, - faceid_embeds=None, - prompt=None, - negative_prompt=None, - scale=1.0, - num_samples=4, - seed=None, - guidance_scale=7.5, - num_inference_steps=30, - s_scale=1.0, - shortcut=True, - **kwargs, + self, + face_image=None, + faceid_embeds=None, + prompt=None, + negative_prompt=None, + scale=1.0, + num_samples=4, + seed=None, + guidance_scale=7.5, + num_inference_steps=30, + s_scale=1.0, + shortcut=True, + **kwargs, ): self.set_scale(scale) @@ -628,20 +786,30 @@ def generate( if prompt is None: prompt = "best quality, high quality" if negative_prompt is None: - negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" + negative_prompt = ( + "monochrome, lowres, bad anatomy, worst quality, low quality" + ) if not isinstance(prompt, List): prompt = [prompt] * num_prompts if not isinstance(negative_prompt, List): negative_prompt = [negative_prompt] * num_prompts - image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds(faceid_embeds, face_image, s_scale, shortcut) + image_prompt_embeds, uncond_image_prompt_embeds = self.get_image_embeds( + faceid_embeds, face_image, s_scale, shortcut + ) bs_embed, seq_len, _ = image_prompt_embeds.shape image_prompt_embeds = image_prompt_embeds.repeat(1, num_samples, 1) - image_prompt_embeds = image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat(1, num_samples, 1) - uncond_image_prompt_embeds = uncond_image_prompt_embeds.view(bs_embed * num_samples, seq_len, -1) + image_prompt_embeds = image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.repeat( + 1, num_samples, 1 + ) + uncond_image_prompt_embeds = uncond_image_prompt_embeds.view( + bs_embed * num_samples, seq_len, -1 + ) with torch.inference_mode(): ( @@ -656,9 +824,13 @@ def generate( negative_prompt=negative_prompt, ) prompt_embeds = torch.cat([prompt_embeds, image_prompt_embeds], dim=1) - negative_prompt_embeds = torch.cat([negative_prompt_embeds, uncond_image_prompt_embeds], dim=1) + negative_prompt_embeds = torch.cat( + [negative_prompt_embeds, uncond_image_prompt_embeds], dim=1 + ) - generator = torch.Generator(self.device).manual_seed(seed) if seed is not None else None + generator = ( + torch.Generator(self.device).manual_seed(seed) if seed is not None else None + ) images = self.pipe( prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, diff --git a/garment_seg/network.py b/garment_seg/network.py index 496a7ed..fe7a890 100644 --- a/garment_seg/network.py +++ b/garment_seg/network.py @@ -14,7 +14,6 @@ def __init__(self, in_ch=3, out_ch=3, dirate=1): self.relu_s1 = nn.ReLU(inplace=True) def forward(self, x): - hx = x xout = self.relu_s1(self.bn_s1(self.conv_s1(hx))) @@ -23,7 +22,6 @@ def forward(self, x): ## upsample tensor 'src' to have the same spatial size with tensor 'tar' def _upsample_like(src, tar): - src = F.upsample(src, size=tar.shape[2:], mode="bilinear") return src @@ -63,7 +61,6 @@ def __init__(self, in_ch=3, mid_ch=12, out_ch=3): self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) def forward(self, x): - hx = x hxin = self.rebnconvin(hx) @@ -142,7 +139,6 @@ def __init__(self, in_ch=3, mid_ch=12, out_ch=3): self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) def forward(self, x): - hx = x hxin = self.rebnconvin(hx) @@ -212,7 +208,6 @@ def __init__(self, in_ch=3, mid_ch=12, out_ch=3): self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) def forward(self, x): - hx = x hxin = self.rebnconvin(hx) @@ -272,7 +267,6 @@ def __init__(self, in_ch=3, mid_ch=12, out_ch=3): self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) def forward(self, x): - hx = x hxin = self.rebnconvin(hx) @@ -322,7 +316,6 @@ def __init__(self, in_ch=3, mid_ch=12, out_ch=3): self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1) def forward(self, x): - hx = x hxin = self.rebnconvin(hx) @@ -384,7 +377,6 @@ def __init__(self, in_ch=3, out_ch=1): self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1) def forward(self, x): - hx = x # stage 1 @@ -494,7 +486,6 @@ def __init__(self, in_ch=3, out_ch=1): self.outconv = nn.Conv2d(6 * out_ch, out_ch, 1) def forward(self, x): - hx = x # stage 1 @@ -556,5 +547,4 @@ def forward(self, x): d0 = self.outconv(torch.cat((d1, d2, d3, d4, d5, d6), 1)) - - return d0, d1, d2, d3, d4, d5, d6 \ No newline at end of file + return d0, d1, d2, d3, d4, d5, d6 diff --git a/garment_seg/process.py b/garment_seg/process.py index 6ab7725..5e10975 100644 --- a/garment_seg/process.py +++ b/garment_seg/process.py @@ -1,4 +1,3 @@ - from .network import U2NET import os @@ -69,7 +68,7 @@ def apply_transform(img): return transform_rgb(img) -def generate_mask(input_image, net, device='cpu'): +def generate_mask(input_image, net, device="cpu"): img = input_image img_size = img.size img = img.resize((768, 768), Image.BICUBIC) @@ -84,13 +83,13 @@ def generate_mask(input_image, net, device='cpu'): output_arr = output_tensor.cpu().numpy() mask = (output_arr != 0).astype(np.uint8) * 255 mask = mask[0] # Selecting the first channel to make it 2D - alpha_mask_img = Image.fromarray(mask, mode='L') + alpha_mask_img = Image.fromarray(mask, mode="L") alpha_mask_img = alpha_mask_img.resize(img_size, Image.BICUBIC) return alpha_mask_img -def load_seg_model(checkpoint_path, device='cpu'): +def load_seg_model(checkpoint_path, device="cpu"): net = U2NET(in_ch=3, out_ch=4) net = load_checkpoint(net, checkpoint_path) net = net.to(device) diff --git a/gradio_animatediff.py b/gradio_animatediff.py index 423782a..d815fa6 100644 --- a/gradio_animatediff.py +++ b/gradio_animatediff.py @@ -2,7 +2,15 @@ import pdb import torch -from diffusers import UniPCMultistepScheduler, AutoencoderKL, DDIMScheduler, MotionAdapter, EulerAncestralDiscreteScheduler, LMSDiscreteScheduler,StableVideoDiffusionPipeline +from diffusers import ( + UniPCMultistepScheduler, + AutoencoderKL, + DDIMScheduler, + MotionAdapter, + EulerAncestralDiscreteScheduler, + LMSDiscreteScheduler, + StableVideoDiffusionPipeline, +) from diffusers.pipelines import AnimateDiffPipeline from PIL import Image import argparse @@ -11,12 +19,13 @@ from pipelines.OmsAnimateDiffusionPipeline import OmsAnimateDiffusionPipeline if __name__ == "__main__": - - parser = argparse.ArgumentParser(description='oms diffusion') - parser.add_argument('--cloth_path', type=str, required=True) - parser.add_argument('--model_path', type=str, required=True) - parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") - parser.add_argument('--output_path', type=str, default="./output_img") + parser = argparse.ArgumentParser(description="oms diffusion") + parser.add_argument("--cloth_path", type=str, required=True) + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" + ) + parser.add_argument("--output_path", type=str, default="./output_img") args = parser.parse_args() @@ -27,12 +36,20 @@ cloth_image = Image.open(args.cloth_path).convert("RGB") - vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) - adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16) + vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to( + dtype=torch.float16 + ) + adapter = MotionAdapter.from_pretrained( + "guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16 + ) - pipe = OmsAnimateDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, motion_adapter=adapter, torch_dtype=torch.float16) + pipe = OmsAnimateDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, motion_adapter=adapter, torch_dtype=torch.float16 + ) # pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter_AnimateDiff(pipe, args.pipe_path, args.model_path, device) - frames, cloth_mask_image = full_net.generate(cloth_image, num_images_per_prompt=1, seed=6896868) + frames, cloth_mask_image = full_net.generate( + cloth_image, num_images_per_prompt=1, seed=6896868 + ) export_to_gif(frames[0], "animation0.gif") diff --git a/gradio_controlnet_inpainting.py b/gradio_controlnet_inpainting.py index 43cec67..166db08 100644 --- a/gradio_controlnet_inpainting.py +++ b/gradio_controlnet_inpainting.py @@ -8,21 +8,35 @@ from pipelines.OmsDiffusionControlNetPipeline import OmsDiffusionControlNetPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--enable_cloth_guidance', action="store_true") -parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument("--enable_cloth_guidance", action="store_true") +parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" +) args = parser.parse_args() device = "cuda" -control_net_openpose = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16) +control_net_openpose = ControlNetModel.from_pretrained( + "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16 +) vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) if args.enable_cloth_guidance: - pipe = OmsDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = OmsDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) else: - pipe = StableDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = StableDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter(pipe, args.model_path, device, args.enable_cloth_guidance) @@ -30,47 +44,130 @@ def make_inpaint_condition(image, image_mask): image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0 - assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size" + assert ( + image.shape[0:1] == image_mask.shape[0:1] + ), "image and image_mask must have the same image size" image[image_mask > 0.5] = -1.0 # set as masked pixel image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) image = torch.from_numpy(image) return image -def process(cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, scale, cloth_guidance_scale, seed, person_image, person_mask): +def process( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + scale, + cloth_guidance_scale, + seed, + person_image, + person_mask, +): inpaint_image = make_inpaint_condition(person_image, person_mask) - images, cloth_mask_image = full_net.generate(cloth_image, cloth_mask_image, prompt, a_prompt, num_samples, n_prompt, seed, scale,cloth_guidance_scale, sample_steps, height, width, image=inpaint_image) + images, cloth_mask_image = full_net.generate( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + num_samples, + n_prompt, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + image=inpaint_image, + ) return images, cloth_mask_image block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") - prompt = gr.Textbox(label="Prompt", value='a photography of a model') + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) + prompt = gr.Textbox(label="Prompt", value="a photography of a model") run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=768, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=576, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) - guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10., value=5. if args.enable_cloth_guidance else 2.5, step=0.1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=2.5, step=0.1, visible=args.enable_cloth_guidance) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) - a_prompt = gr.Textbox(label="Added Prompt", value='best quality, high quality') - n_prompt = gr.Textbox(label="Negative Prompt", value='bare, monochrome, lowres, bad anatomy, worst quality, low quality') + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=768, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=576, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=10.0, + value=5.0 if args.enable_cloth_guidance else 2.5, + step=0.1, + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=2.5, + step=0.1, + visible=args.enable_cloth_guidance, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) + a_prompt = gr.Textbox( + label="Added Prompt", value="best quality, high quality" + ) + n_prompt = gr.Textbox( + label="Negative Prompt", + value="bare, monochrome, lowres, bad anatomy, worst quality, low quality", + ) with gr.Column(): person_image = gr.Image(label="person Image", type="pil") person_mask = gr.Image(label="person mask", type="pil") with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", min_width=384) - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery", min_width=384 + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) - ips = [cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, guidance_scale, cloth_guidance_scale, seed, person_image, person_mask] + ips = [ + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + guidance_scale, + cloth_guidance_scale, + seed, + person_image, + person_mask, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) block.launch(server_name="0.0.0.0", server_port=7860) diff --git a/gradio_controlnet_openpose.py b/gradio_controlnet_openpose.py index 5428d68..845d04d 100644 --- a/gradio_controlnet_openpose.py +++ b/gradio_controlnet_openpose.py @@ -7,22 +7,36 @@ from garment_adapter.garment_diffusion import ClothAdapter from pipelines.OmsDiffusionControlNetPipeline import OmsDiffusionControlNetPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--enable_cloth_guidance', action="store_true") -parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument("--enable_cloth_guidance", action="store_true") +parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" +) args = parser.parse_args() device = "cuda" openpose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet").to(device) -control_net_openpose = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16) +control_net_openpose = ControlNetModel.from_pretrained( + "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16 +) vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) if args.enable_cloth_guidance: - pipe = OmsDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = OmsDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) else: - pipe = StableDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = StableDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter(pipe, args.model_path, device, args.enable_cloth_guidance) @@ -32,40 +46,119 @@ def get_pose(image): return openpose_image -def process(cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, scale, cloth_guidance_scale, seed, pose_image): - images, cloth_mask_image = full_net.generate(cloth_image, cloth_mask_image, prompt, a_prompt, num_samples, n_prompt, seed, scale, cloth_guidance_scale, sample_steps, height, width, image=pose_image) +def process( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + scale, + cloth_guidance_scale, + seed, + pose_image, +): + images, cloth_mask_image = full_net.generate( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + num_samples, + n_prompt, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + image=pose_image, + ) return images, cloth_mask_image block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") - prompt = gr.Textbox(label="Prompt", value='a photography of a model') + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) + prompt = gr.Textbox(label="Prompt", value="a photography of a model") run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=768, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=576, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=768, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=576, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) - guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10., value=5. if args.enable_cloth_guidance else 2.5, step=0.1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=2.5, step=0.1, visible=args.enable_cloth_guidance) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) - a_prompt = gr.Textbox(label="Added Prompt", value='best quality, high quality') - n_prompt = gr.Textbox(label="Negative Prompt", value='bare, monochrome, lowres, bad anatomy, worst quality, low quality') + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=10.0, + value=5.0 if args.enable_cloth_guidance else 2.5, + step=0.1, + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=2.5, + step=0.1, + visible=args.enable_cloth_guidance, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) + a_prompt = gr.Textbox( + label="Added Prompt", value="best quality, high quality" + ) + n_prompt = gr.Textbox( + label="Negative Prompt", + value="bare, monochrome, lowres, bad anatomy, worst quality, low quality", + ) with gr.Column(): pose_image = gr.Image(label="pose Image", type="pil") pose_button = gr.Button(value="get pose") with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery", min_width=384) - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery", min_width=384 + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) - ips = [cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, guidance_scale, cloth_guidance_scale, seed, pose_image] + ips = [ + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + guidance_scale, + cloth_guidance_scale, + seed, + pose_image, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) pose_button.click(fn=get_pose, inputs=pose_image, outputs=pose_image) diff --git a/gradio_generate.py b/gradio_generate.py index 0149c38..743105b 100644 --- a/gradio_generate.py +++ b/gradio_generate.py @@ -7,10 +7,12 @@ from garment_adapter.garment_diffusion import ClothAdapter from pipelines.OmsDiffusionPipeline import OmsDiffusionPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--enable_cloth_guidance', action="store_true") -parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument("--enable_cloth_guidance", action="store_true") +parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" +) args = parser.parse_args() @@ -18,44 +20,124 @@ vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) if args.enable_cloth_guidance: - pipe = OmsDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = OmsDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) else: - pipe = StableDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = StableDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter(pipe, args.model_path, device, args.enable_cloth_guidance) -def process(cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, scale, cloth_guidance_scale, seed): - images, cloth_mask_image = full_net.generate(cloth_image, cloth_mask_image, prompt, a_prompt, num_samples, n_prompt, seed, scale, cloth_guidance_scale, sample_steps, height, width) +def process( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + scale, + cloth_guidance_scale, + seed, +): + images, cloth_mask_image = full_net.generate( + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + num_samples, + n_prompt, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + ) return images, cloth_mask_image block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") - prompt = gr.Textbox(label="Prompt", value='a photography of a model') + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) + prompt = gr.Textbox(label="Prompt", value="a photography of a model") run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=768, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=576, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) - guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10., value=5. if args.enable_cloth_guidance else 2.5, step=0.1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=2.5, step=0.1, visible=args.enable_cloth_guidance) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) - a_prompt = gr.Textbox(label="Added Prompt", value='best quality, high quality') - n_prompt = gr.Textbox(label="Negative Prompt", value='bare, monochrome, lowres, bad anatomy, worst quality, low quality') + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=768, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=576, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=10.0, + value=5.0 if args.enable_cloth_guidance else 2.5, + step=0.1, + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=2.5, + step=0.1, + visible=args.enable_cloth_guidance, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) + a_prompt = gr.Textbox( + label="Added Prompt", value="best quality, high quality" + ) + n_prompt = gr.Textbox( + label="Negative Prompt", + value="bare, monochrome, lowres, bad anatomy, worst quality, low quality", + ) with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery") - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery" + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) - ips = [cloth_image, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, guidance_scale, cloth_guidance_scale, seed] + ips = [ + cloth_image, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + guidance_scale, + cloth_guidance_scale, + seed, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) block.launch(server_name="0.0.0.0", server_port=7860) diff --git a/gradio_ipadapter_faceid.py b/gradio_ipadapter_faceid.py index 429aeae..2bdcdf0 100644 --- a/gradio_ipadapter_faceid.py +++ b/gradio_ipadapter_faceid.py @@ -10,11 +10,18 @@ from pipelines.OmsDiffusionPipeline import OmsDiffusionPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") -parser.add_argument('--enable_cloth_guidance', action="store_true") -parser.add_argument('--faceid_version', type=str, default="FaceIDPlusV2", choices=['FaceID', 'FaceIDPlus', 'FaceIDPlusV2']) +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" +) +parser.add_argument("--enable_cloth_guidance", action="store_true") +parser.add_argument( + "--faceid_version", + type=str, + default="FaceIDPlusV2", + choices=["FaceID", "FaceIDPlus", "FaceIDPlusV2"], +) args = parser.parse_args() @@ -22,9 +29,13 @@ vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) if args.enable_cloth_guidance: - pipe = OmsDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = OmsDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) else: - pipe = StableDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = StableDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) if args.faceid_version == "FaceID": @@ -34,7 +45,9 @@ pipe.fuse_lora() from garment_adapter.garment_ipadapter_faceid import IPAdapterFaceID - ip_model = IPAdapterFaceID(pipe, args.model_path, ip_ckpt, device, args.enable_cloth_guidance) + ip_model = IPAdapterFaceID( + pipe, args.model_path, ip_ckpt, device, args.enable_cloth_guidance + ) else: if args.faceid_version == "FaceIDPlus": ip_ckpt = "./checkpoints/ipadapter_faceid/ip-adapter-faceid-plus_sd15.bin" @@ -48,16 +61,68 @@ pipe.load_lora_weights(ip_lora) pipe.fuse_lora() image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" - from garment_adapter.garment_ipadapter_faceid import IPAdapterFaceIDPlus as IPAdapterFaceID - - ip_model = IPAdapterFaceID(pipe, args.model_path, image_encoder_path, ip_ckpt, device, args.enable_cloth_guidance) - - -def process(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, scale, cloth_guidance_scale, seed): + from garment_adapter.garment_ipadapter_faceid import ( + IPAdapterFaceIDPlus as IPAdapterFaceID, + ) + + ip_model = IPAdapterFaceID( + pipe, + args.model_path, + image_encoder_path, + ip_ckpt, + device, + args.enable_cloth_guidance, + ) + + +def process( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + scale, + cloth_guidance_scale, + seed, +): if args.faceid_version == "FaceID": - result = ip_model.generate(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, seed, scale, cloth_guidance_scale, sample_steps, height, width) + result = ip_model.generate( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + ) else: - result = ip_model.generate(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, seed, scale, cloth_guidance_scale, sample_steps, height, width, shortcut=v2) + result = ip_model.generate( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + shortcut=v2, + ) if result is None: raise gr.Error("人脸检测异常，尝试其他肖像") else: @@ -68,30 +133,81 @@ def process(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): face_img = gr.Image(label="face Image", type="pil") cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") - prompt = gr.Textbox(label="Prompt", value='a photography') + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) + prompt = gr.Textbox(label="Prompt", value="a photography") run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=768, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=576, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) - guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10., value=3. if args.enable_cloth_guidance else 2.5, step=0.1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=3., step=0.1, visible=args.enable_cloth_guidance) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) - a_prompt = gr.Textbox(label="Added Prompt", value='best quality, high quality') - n_prompt = gr.Textbox(label="Negative Prompt", value='bare, monochrome, lowres, bad anatomy, worst quality, low quality') + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=768, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=576, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=10.0, + value=3.0 if args.enable_cloth_guidance else 2.5, + step=0.1, + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=3.0, + step=0.1, + visible=args.enable_cloth_guidance, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) + a_prompt = gr.Textbox( + label="Added Prompt", value="best quality, high quality" + ) + n_prompt = gr.Textbox( + label="Negative Prompt", + value="bare, monochrome, lowres, bad anatomy, worst quality, low quality", + ) with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery") - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) - - ips = [cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, guidance_scale, cloth_guidance_scale, seed] + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery" + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) + + ips = [ + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + guidance_scale, + cloth_guidance_scale, + seed, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) block.launch(server_name="0.0.0.0", server_port=7860) diff --git a/gradio_ipadapter_openpose.py b/gradio_ipadapter_openpose.py index 3401b40..98c563e 100644 --- a/gradio_ipadapter_openpose.py +++ b/gradio_ipadapter_openpose.py @@ -10,23 +10,42 @@ from pipelines.OmsDiffusionControlNetPipeline import OmsDiffusionControlNetPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") -parser.add_argument('--enable_cloth_guidance', action="store_true") -parser.add_argument('--faceid_version', type=str, default="FaceIDPlus", choices=['FaceID', 'FaceIDPlus', 'FaceIDPlusV2']) +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" +) +parser.add_argument("--enable_cloth_guidance", action="store_true") +parser.add_argument( + "--faceid_version", + type=str, + default="FaceIDPlus", + choices=["FaceID", "FaceIDPlus", "FaceIDPlusV2"], +) args = parser.parse_args() device = "cuda" openpose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet").to(device) -control_net_openpose = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16) +control_net_openpose = ControlNetModel.from_pretrained( + "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16 +) vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) if args.enable_cloth_guidance: - pipe = OmsDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = OmsDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) else: - pipe = StableDiffusionControlNetPipeline.from_pretrained(args.pipe_path, vae=vae, controlnet=control_net_openpose, torch_dtype=torch.float16) + pipe = StableDiffusionControlNetPipeline.from_pretrained( + args.pipe_path, + vae=vae, + controlnet=control_net_openpose, + torch_dtype=torch.float16, + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) if args.faceid_version == "FaceID": @@ -37,7 +56,9 @@ pipe.fuse_lora() from garment_adapter.garment_ipadapter_faceid import IPAdapterFaceID - ip_model = IPAdapterFaceID(pipe, args.model_path, ip_ckpt, device, args.enable_cloth_guidance) + ip_model = IPAdapterFaceID( + pipe, args.model_path, ip_ckpt, device, args.enable_cloth_guidance + ) else: if args.faceid_version == "FaceIDPlus": ip_ckpt = "./checkpoints/ipadapter_faceid/ip-adapter-faceid-plus_sd15.bin" @@ -52,16 +73,71 @@ pipe.fuse_lora() image_encoder_path = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" - from garment_adapter.garment_ipadapter_faceid import IPAdapterFaceIDPlus as IPAdapterFaceID - - ip_model = IPAdapterFaceID(pipe, args.model_path, image_encoder_path, ip_ckpt, device, args.enable_cloth_guidance) - - -def process(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, scale, cloth_guidance_scale, seed, pose_image): + from garment_adapter.garment_ipadapter_faceid import ( + IPAdapterFaceIDPlus as IPAdapterFaceID, + ) + + ip_model = IPAdapterFaceID( + pipe, + args.model_path, + image_encoder_path, + ip_ckpt, + device, + args.enable_cloth_guidance, + ) + + +def process( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + scale, + cloth_guidance_scale, + seed, + pose_image, +): if args.faceid_version == "FaceID": - result = ip_model.generate(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, seed, scale, cloth_guidance_scale, sample_steps, height, width, image=pose_image) + result = ip_model.generate( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + image=pose_image, + ) else: - result = ip_model.generate(cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, seed, scale, cloth_guidance_scale, sample_steps, height, width, shortcut=v2, image=pose_image) + result = ip_model.generate( + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + seed, + scale, + cloth_guidance_scale, + sample_steps, + height, + width, + shortcut=v2, + image=pose_image, + ) if result is None: raise gr.Error("人脸检测异常，尝试其他肖像") else: @@ -77,32 +153,84 @@ def get_pose(image): block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): face_img = gr.Image(label="face Image", type="pil") cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") - prompt = gr.Textbox(label="Prompt", value='a photography') + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) + prompt = gr.Textbox(label="Prompt", value="a photography") run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=768, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=576, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) - guidance_scale = gr.Slider(label="Guidance Scale", minimum=1, maximum=10., value=3. if args.enable_cloth_guidance else 2.5, step=0.1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=3., step=0.1, visible=args.enable_cloth_guidance) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) - a_prompt = gr.Textbox(label="Added Prompt", value='best quality, high quality') - n_prompt = gr.Textbox(label="Negative Prompt", value='bare, monochrome, lowres, bad anatomy, worst quality, low quality') + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=768, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=576, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) + guidance_scale = gr.Slider( + label="Guidance Scale", + minimum=1, + maximum=10.0, + value=3.0 if args.enable_cloth_guidance else 2.5, + step=0.1, + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=3.0, + step=0.1, + visible=args.enable_cloth_guidance, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) + a_prompt = gr.Textbox( + label="Added Prompt", value="best quality, high quality" + ) + n_prompt = gr.Textbox( + label="Negative Prompt", + value="bare, monochrome, lowres, bad anatomy, worst quality, low quality", + ) with gr.Column(): pose_image = gr.Image(label="pose Image", type="pil") pose_button = gr.Button(value="get pose") with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery") - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) - - ips = [cloth_image, face_img, cloth_mask_image, prompt, a_prompt, n_prompt, num_samples, width, height, sample_steps, guidance_scale, cloth_guidance_scale, seed, pose_image] + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery" + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) + + ips = [ + cloth_image, + face_img, + cloth_mask_image, + prompt, + a_prompt, + n_prompt, + num_samples, + width, + height, + sample_steps, + guidance_scale, + cloth_guidance_scale, + seed, + pose_image, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) pose_button.click(fn=get_pose, inputs=pose_image, outputs=pose_image) diff --git a/gradio_sd_inpainting.py b/gradio_sd_inpainting.py index eee6377..c3aba3c 100644 --- a/gradio_sd_inpainting.py +++ b/gradio_sd_inpainting.py @@ -9,54 +9,116 @@ from garment_adapter.garment_diffusion import ClothAdapter from pipelines.OmsDiffusionInpaintPipeline import OmsDiffusionInpaintPipeline -parser = argparse.ArgumentParser(description='oms diffusion') -parser.add_argument('--model_path', type=str, required=True) -parser.add_argument('--pipe_path', type=str, default="runwayml/stable-diffusion-inpainting") +parser = argparse.ArgumentParser(description="oms diffusion") +parser.add_argument("--model_path", type=str, required=True) +parser.add_argument( + "--pipe_path", type=str, default="runwayml/stable-diffusion-inpainting" +) args = parser.parse_args() device = "cuda" vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) -pipe = OmsDiffusionInpaintPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) +pipe = OmsDiffusionInpaintPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 +) pipe.safety_checker = None pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter(pipe, args.model_path, device, False) -def process(person_image, person_mask, cloth_image, cloth_mask_image, num_samples, width, height, sample_steps, cloth_guidance_scale, seed): +def process( + person_image, + person_mask, + cloth_image, + cloth_mask_image, + num_samples, + width, + height, + sample_steps, + cloth_guidance_scale, + seed, +): # person_image = person_image_mask['background'].convert("RGB") # person_mask = person_image_mask['layers'][0].split()[-1] - images, cloth_mask_image = full_net.generate_inpainting(cloth_image, cloth_mask_image, num_samples, seed, cloth_guidance_scale, sample_steps, height, width, image=person_image, mask_image=person_mask) + images, cloth_mask_image = full_net.generate_inpainting( + cloth_image, + cloth_mask_image, + num_samples, + seed, + cloth_guidance_scale, + sample_steps, + height, + width, + image=person_image, + mask_image=person_mask, + ) return images, cloth_mask_image block = gr.Blocks().queue() with block: with gr.Row(): - gr.Markdown("##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##") + gr.Markdown( + "##You can enlarge image resolution to get better face, but the cloth maybe lose control, we will release high-resolution checkpoint soon##" + ) with gr.Row(): with gr.Column(): cloth_image = gr.Image(label="cloth Image", type="pil") - cloth_mask_image = gr.Image(label="cloth mask Image, if not support, will be produced by inner segment algorithm", type="pil") + cloth_mask_image = gr.Image( + label="cloth mask Image, if not support, will be produced by inner segment algorithm", + type="pil", + ) run_button = gr.Button(value="Run") with gr.Accordion("Advanced options", open=False): - num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1) - height = gr.Slider(label="Height", minimum=256, maximum=1024, value=1024, step=64) - width = gr.Slider(label="Width", minimum=192, maximum=768, value=768, step=64) - sample_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1) - cloth_guidance_scale = gr.Slider(label="Cloth guidance Scale", minimum=1, maximum=10., value=2.5, step=0.1) - seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234) + num_samples = gr.Slider( + label="Images", minimum=1, maximum=12, value=1, step=1 + ) + height = gr.Slider( + label="Height", minimum=256, maximum=1024, value=1024, step=64 + ) + width = gr.Slider( + label="Width", minimum=192, maximum=768, value=768, step=64 + ) + sample_steps = gr.Slider( + label="Steps", minimum=1, maximum=100, value=20, step=1 + ) + cloth_guidance_scale = gr.Slider( + label="Cloth guidance Scale", + minimum=1, + maximum=10.0, + value=2.5, + step=0.1, + ) + seed = gr.Slider( + label="Seed", minimum=-1, maximum=2147483647, step=1, value=1234 + ) with gr.Column(): person_image = gr.Image(label="person Image", type="pil") person_mask = gr.Image(label="person mask", type="pil") # person_image_mask = gr.ImageMask(label="person Image", type="pil") with gr.Column(): - result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery") - cloth_seg_image = gr.Image(label="cloth mask", type="pil", width=192, height=256) + result_gallery = gr.Gallery( + label="Output", show_label=False, elem_id="gallery" + ) + cloth_seg_image = gr.Image( + label="cloth mask", type="pil", width=192, height=256 + ) - ips = [person_image, person_mask, cloth_image, cloth_mask_image, num_samples, width, height, sample_steps, cloth_guidance_scale, seed] + ips = [ + person_image, + person_mask, + cloth_image, + cloth_mask_image, + num_samples, + width, + height, + sample_steps, + cloth_guidance_scale, + seed, + ] run_button.click(fn=process, inputs=ips, outputs=[result_gallery, cloth_seg_image]) block.launch(server_name="0.0.0.0", server_port=7860) diff --git a/inference.py b/inference.py index 3c5d5af..e82525b 100644 --- a/inference.py +++ b/inference.py @@ -11,13 +11,14 @@ from pipelines.OmsDiffusionPipeline import OmsDiffusionPipeline if __name__ == "__main__": - - parser = argparse.ArgumentParser(description='oms diffusion') - parser.add_argument('--cloth_path', type=str, required=True) - parser.add_argument('--model_path', type=str, required=True) - parser.add_argument('--enable_cloth_guidance', action="store_true") - parser.add_argument('--pipe_path', type=str, default="SG161222/Realistic_Vision_V4.0_noVAE") - parser.add_argument('--output_path', type=str, default="./output_img") + parser = argparse.ArgumentParser(description="oms diffusion") + parser.add_argument("--cloth_path", type=str, required=True) + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--enable_cloth_guidance", action="store_true") + parser.add_argument( + "--pipe_path", type=str, default="SG161222/Realistic_Vision_V4.0_noVAE" + ) + parser.add_argument("--output_path", type=str, default="./output_img") args = parser.parse_args() @@ -28,11 +29,17 @@ cloth_image = Image.open(args.cloth_path).convert("RGB") - vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(dtype=torch.float16) + vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to( + dtype=torch.float16 + ) if args.enable_cloth_guidance: - pipe = OmsDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = OmsDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) else: - pipe = StableDiffusionPipeline.from_pretrained(args.pipe_path, vae=vae, torch_dtype=torch.float16) + pipe = StableDiffusionPipeline.from_pretrained( + args.pipe_path, vae=vae, torch_dtype=torch.float16 + ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) full_net = ClothAdapter(pipe, args.model_path, device, args.enable_cloth_guidance) diff --git a/pipelines/OmsAnimateDiffusionPipeline.py b/pipelines/OmsAnimateDiffusionPipeline.py index 8c94f9e..b6cfcb6 100644 --- a/pipelines/OmsAnimateDiffusionPipeline.py +++ b/pipelines/OmsAnimateDiffusionPipeline.py @@ -1,7 +1,7 @@ from diffusers.pipelines.animatediff.pipeline_animatediff import * -class OmsAnimateDiffusionPipeline(AnimateDiffPipeline): +class OmsAnimateDiffusionPipeline(AnimateDiffPipeline): def _denoise_loop( self, timesteps, @@ -25,8 +25,12 @@ def _denoise_loop( with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 3) if do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = ( + torch.cat([latents] * 3) if do_classifier_free_guidance else latents + ) + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t + ) # predict the noise residual noise_pred = self.unet( @@ -39,15 +43,19 @@ def _denoise_loop( # perform guidance if do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cloth, noise_pred_text = noise_pred.chunk(3) + noise_pred_uncond, noise_pred_cloth, noise_pred_text = ( + noise_pred.chunk(3) + ) noise_pred = ( - noise_pred_uncond - + guidance_scale * (noise_pred_text - noise_pred_cloth) - + cloth_guidance_scale * (noise_pred_cloth - noise_pred_uncond) + noise_pred_uncond + + guidance_scale * (noise_pred_text - noise_pred_cloth) + + cloth_guidance_scale * (noise_pred_cloth - noise_pred_uncond) ) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample + latents = self.scheduler.step( + noise_pred, t, latents, **extra_step_kwargs + ).prev_sample if callback_on_step_end is not None: callback_kwargs = {} @@ -57,10 +65,14 @@ def _denoise_loop( latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): progress_bar.update() if callback is not None and i % callback_steps == 0: callback(i, t, latents) @@ -69,29 +81,29 @@ def _denoise_loop( @torch.no_grad() def __call__( - self, - prompt: Union[str, List[str]] = None, - num_frames: Optional[int] = 16, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - guidance_scale: float = 7.5, - cloth_guidance_scale: float = 7.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_videos_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - **kwargs, + self, + prompt: Union[str, List[str]] = None, + num_frames: Optional[int] = 16, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + cloth_guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_videos_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, ): r""" The call function to the pipeline for generation. @@ -214,7 +226,9 @@ def __call__( # 3. Encode input prompt text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None + else None ) prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, @@ -231,7 +245,9 @@ def __call__( # Here we concatenate the unconditional and text embeddings into a single batch # to avoid doing two forward passes if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, negative_prompt_embeds, prompt_embeds]) + prompt_embeds = torch.cat( + [negative_prompt_embeds, negative_prompt_embeds, prompt_embeds] + ) if ip_adapter_image is not None: image_embeds = self.prepare_ip_adapter_image_embeds( @@ -261,7 +277,9 @@ def __call__( extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 7. Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None + added_cond_kwargs = ( + {"image_embeds": image_embeds} if ip_adapter_image is not None else None + ) # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order @@ -303,4 +321,4 @@ def __call__( # 9. Offload all models self.maybe_free_model_hooks() - return video \ No newline at end of file + return video diff --git a/pipelines/OmsDiffusionControlNetPipeline.py b/pipelines/OmsDiffusionControlNetPipeline.py index 0e6c2e2..1e13e9d 100644 --- a/pipelines/OmsDiffusionControlNetPipeline.py +++ b/pipelines/OmsDiffusionControlNetPipeline.py @@ -5,34 +5,34 @@ class OmsDiffusionControlNetPipeline(StableDiffusionControlNetPipeline): @torch.no_grad() def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 7.5, - cloth_guidance_scale: float = 2.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - controlnet_conditioning_scale: Union[float, List[float]] = 1.0, - guess_mode: bool = False, - control_guidance_start: Union[float, List[float]] = 0.0, - control_guidance_end: Union[float, List[float]] = 1.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - **kwargs, + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + timesteps: List[int] = None, + guidance_scale: float = 7.5, + cloth_guidance_scale: float = 2.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + controlnet_conditioning_scale: Union[float, List[float]] = 1.0, + guess_mode: bool = False, + control_guidance_start: Union[float, List[float]] = 0.0, + control_guidance_end: Union[float, List[float]] = 1.0, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, ): r""" The call function to the pipeline for generation. @@ -150,15 +150,31 @@ def __call__( "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`", ) - controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet + controlnet = ( + self.controlnet._orig_mod + if is_compiled_module(self.controlnet) + else self.controlnet + ) # align format for control guidance - if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list): - control_guidance_start = len(control_guidance_end) * [control_guidance_start] - elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list): + if not isinstance(control_guidance_start, list) and isinstance( + control_guidance_end, list + ): + control_guidance_start = len(control_guidance_end) * [ + control_guidance_start + ] + elif not isinstance(control_guidance_end, list) and isinstance( + control_guidance_start, list + ): control_guidance_end = len(control_guidance_start) * [control_guidance_end] - elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list): - mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1 + elif not isinstance(control_guidance_start, list) and not isinstance( + control_guidance_end, list + ): + mult = ( + len(controlnet.nets) + if isinstance(controlnet, MultiControlNetModel) + else 1 + ) control_guidance_start, control_guidance_end = ( mult * [control_guidance_start], mult * [control_guidance_end], @@ -192,8 +208,12 @@ def __call__( device = self._execution_device - if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float): - controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets) + if isinstance(controlnet, MultiControlNetModel) and isinstance( + controlnet_conditioning_scale, float + ): + controlnet_conditioning_scale = [controlnet_conditioning_scale] * len( + controlnet.nets + ) global_pool_conditions = ( controlnet.config.global_pool_conditions @@ -204,7 +224,9 @@ def __call__( # 3. Encode input prompt text_encoder_lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None + else None ) prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, @@ -221,7 +243,9 @@ def __call__( # Here we concatenate the unconditional and text embeddings into a single batch # to avoid doing two forward passes if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, negative_prompt_embeds, prompt_embeds]) + prompt_embeds = torch.cat( + [negative_prompt_embeds, negative_prompt_embeds, prompt_embeds] + ) if ip_adapter_image is not None: image_embeds = self.prepare_ip_adapter_image_embeds( @@ -243,7 +267,7 @@ def __call__( ) if self.do_classifier_free_guidance and not guess_mode: image = image.chunk(2)[0] - image = torch.cat([image]*3) + image = torch.cat([image] * 3) height, width = image.shape[-2:] elif isinstance(controlnet, MultiControlNetModel): images = [] @@ -274,7 +298,9 @@ def __call__( assert False # 5. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps + ) self._num_timesteps = len(timesteps) # 6. Prepare latent variables @@ -293,7 +319,9 @@ def __call__( # 6.5 Optionally get Guidance Scale Embedding timestep_cond = None if self.unet.config.time_cond_proj_dim is not None: - guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat( + batch_size * num_images_per_prompt + ) timestep_cond = self.get_guidance_scale_embedding( guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim ).to(device=device, dtype=latents.dtype) @@ -302,7 +330,9 @@ def __call__( extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 7.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None + added_cond_kwargs = ( + {"image_embeds": image_embeds} if ip_adapter_image is not None else None + ) # 7.2 Create tensor stating which controlnets to keep controlnet_keep = [] @@ -311,7 +341,9 @@ def __call__( 1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e) for s, e in zip(control_guidance_start, control_guidance_end) ] - controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps) + controlnet_keep.append( + keeps[0] if isinstance(controlnet, ControlNetModel) else keeps + ) # 8. Denoising loop num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order @@ -322,24 +354,39 @@ def __call__( for i, t in enumerate(timesteps): # Relevant thread: # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428 - if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1: + if ( + is_unet_compiled and is_controlnet_compiled + ) and is_torch_higher_equal_2_1: torch._inductor.cudagraph_mark_step_begin() # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = ( + torch.cat([latents] * 3) + if self.do_classifier_free_guidance + else latents + ) + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t + ) # controlnet(s) inference if guess_mode and self.do_classifier_free_guidance: # Infer ControlNet only for the conditional batch. control_model_input = latents - control_model_input = self.scheduler.scale_model_input(control_model_input, t) + control_model_input = self.scheduler.scale_model_input( + control_model_input, t + ) controlnet_prompt_embeds = prompt_embeds.chunk(3)[1] else: control_model_input = latent_model_input controlnet_prompt_embeds = prompt_embeds if isinstance(controlnet_keep[i], list): - cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])] + cond_scale = [ + c * s + for c, s in zip( + controlnet_conditioning_scale, controlnet_keep[i] + ) + ] else: controlnet_cond_scale = controlnet_conditioning_scale if isinstance(controlnet_cond_scale, list): @@ -360,8 +407,13 @@ def __call__( # Infered ControlNet only for the conditional batch. # To apply the output of ControlNet to both the unconditional and conditional batches, # add 0 to the unconditional batch to keep it unchanged. - down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples] - mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample]) + down_block_res_samples = [ + torch.cat([torch.zeros_like(d), d]) + for d in down_block_res_samples + ] + mid_block_res_sample = torch.cat( + [torch.zeros_like(mid_block_res_sample), mid_block_res_sample] + ) # predict the noise residual noise_pred = self.unet( @@ -378,15 +430,19 @@ def __call__( # perform guidance if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cloth, noise_pred_text = noise_pred.chunk(3) + noise_pred_uncond, noise_pred_cloth, noise_pred_text = ( + noise_pred.chunk(3) + ) noise_pred = ( - noise_pred_uncond - + guidance_scale * (noise_pred_text - noise_pred_cloth) - + cloth_guidance_scale * (noise_pred_cloth - noise_pred_uncond) + noise_pred_uncond + + guidance_scale * (noise_pred_text - noise_pred_cloth) + + cloth_guidance_scale * (noise_pred_cloth - noise_pred_uncond) ) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + latents = self.scheduler.step( + noise_pred, t, latents, **extra_step_kwargs, return_dict=False + )[0] if callback_on_step_end is not None: callback_kwargs = {} @@ -396,10 +452,14 @@ def __call__( latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): progress_bar.update() if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) @@ -413,10 +473,14 @@ def __call__( torch.cuda.empty_cache() if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + image = self.vae.decode( + latents / self.vae.config.scaling_factor, + return_dict=False, + generator=generator, + )[0] + image, has_nsfw_concept = self.run_safety_checker( + image, device, prompt_embeds.dtype + ) else: image = latents has_nsfw_concept = None @@ -426,7 +490,9 @@ def __call__( else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + image = self.image_processor.postprocess( + image, output_type=output_type, do_denormalize=do_denormalize + ) # Offload all models self.maybe_free_model_hooks() @@ -434,4 +500,6 @@ def __call__( if not return_dict: return (image, has_nsfw_concept) - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) + return StableDiffusionPipelineOutput( + images=image, nsfw_content_detected=has_nsfw_concept + ) diff --git a/pipelines/OmsDiffusionInpaintPipeline.py b/pipelines/OmsDiffusionInpaintPipeline.py index ec36d7e..5a704fd 100644 --- a/pipelines/OmsDiffusionInpaintPipeline.py +++ b/pipelines/OmsDiffusionInpaintPipeline.py @@ -4,37 +4,36 @@ class OmsDiffusionInpaintPipeline(StableDiffusionInpaintPipeline): - @torch.no_grad() def __call__( - self, - prompt: Union[str, List[str]] = None, - image: PipelineImageInput = None, - mask_image: PipelineImageInput = None, - masked_image_latents: torch.FloatTensor = None, - height: Optional[int] = None, - width: Optional[int] = None, - padding_mask_crop: Optional[int] = None, - strength: float = 1.0, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 0., - cloth_guidance_scale: float = 2.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - clip_skip: int = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - **kwargs, + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + mask_image: PipelineImageInput = None, + masked_image_latents: torch.FloatTensor = None, + height: Optional[int] = None, + width: Optional[int] = None, + padding_mask_crop: Optional[int] = None, + strength: float = 1.0, + num_inference_steps: int = 50, + timesteps: List[int] = None, + guidance_scale: float = 0.0, + cloth_guidance_scale: float = 2.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + clip_skip: int = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, ): r""" The call function to the pipeline for generation. @@ -200,8 +199,8 @@ def __call__( padding_mask_crop, ) - self._guidance_scale = 0. - self.cloth_classifier_free_guidance = cloth_guidance_scale > 1. + self._guidance_scale = 0.0 + self.cloth_classifier_free_guidance = cloth_guidance_scale > 1.0 self._clip_skip = clip_skip self._cross_attention_kwargs = cross_attention_kwargs self._interrupt = False @@ -218,7 +217,9 @@ def __call__( # 3. Encode input prompt text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + cross_attention_kwargs.get("scale", None) + if cross_attention_kwargs is not None + else None ) prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, @@ -243,7 +244,9 @@ def __call__( ) # 4. set timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps + ) timesteps, num_inference_steps = self.get_timesteps( num_inference_steps=num_inference_steps, strength=strength, device=device ) @@ -261,7 +264,9 @@ def __call__( # 5. Preprocess mask and image if padding_mask_crop is not None: - crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop) + crops_coords = self.mask_processor.get_crop_region( + mask_image, width, height, pad=padding_mask_crop + ) resize_mode = "fill" else: crops_coords = None @@ -269,7 +274,11 @@ def __call__( original_image = image init_image = self.image_processor.preprocess( - image, height=height, width=width, crops_coords=crops_coords, resize_mode=resize_mode + image, + height=height, + width=width, + crops_coords=crops_coords, + resize_mode=resize_mode, ) init_image = init_image.to(dtype=torch.float32) @@ -301,7 +310,11 @@ def __call__( # 7. Prepare mask latent variables mask_condition = self.mask_processor.preprocess( - mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords + mask_image, + height=height, + width=width, + resize_mode=resize_mode, + crops_coords=crops_coords, ) if masked_image_latents is None: @@ -326,7 +339,10 @@ def __call__( # default case for runwayml/stable-diffusion-inpainting num_channels_mask = mask.shape[1] num_channels_masked_image = masked_image_latents.shape[1] - if num_channels_latents + num_channels_mask + num_channels_masked_image != self.unet.config.in_channels: + if ( + num_channels_latents + num_channels_mask + num_channels_masked_image + != self.unet.config.in_channels + ): raise ValueError( f"Incorrect configuration settings! The config of `pipeline.unet`: {self.unet.config} expects" f" {self.unet.config.in_channels} but received `num_channels_latents`: {num_channels_latents} +" @@ -343,12 +359,16 @@ def __call__( extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 9.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None + added_cond_kwargs = ( + {"image_embeds": image_embeds} if ip_adapter_image is not None else None + ) # 9.2 Optionally get Guidance Scale Embedding timestep_cond = None if self.unet.config.time_cond_proj_dim is not None: - guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat( + batch_size * num_images_per_prompt + ) timestep_cond = self.get_guidance_scale_embedding( guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim ).to(device=device, dtype=latents.dtype) @@ -362,13 +382,21 @@ def __call__( continue # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.cloth_classifier_free_guidance else latents + latent_model_input = ( + torch.cat([latents] * 2) + if self.cloth_classifier_free_guidance + else latents + ) # concat latents, mask, masked_image_latents in the channel dimension - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t + ) if num_channels_unet == 9: - latent_model_input = torch.cat([latent_model_input, mask, masked_image_latents], dim=1) + latent_model_input = torch.cat( + [latent_model_input, mask, masked_image_latents], dim=1 + ) # predict the noise residual noise_pred = self.unet( @@ -384,10 +412,14 @@ def __call__( # perform guidance if self.cloth_classifier_free_guidance: noise_pred_uncond, noise_pred_cloth = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + cloth_guidance_scale * (noise_pred_cloth - noise_pred_uncond) + noise_pred = noise_pred_uncond + cloth_guidance_scale * ( + noise_pred_cloth - noise_pred_uncond + ) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + latents = self.scheduler.step( + noise_pred, t, latents, **extra_step_kwargs, return_dict=False + )[0] if num_channels_unet == 4: init_latents_proper = image_latents if self.cloth_classifier_free_guidance: @@ -401,7 +433,9 @@ def __call__( init_latents_proper, noise, torch.tensor([noise_timestep]) ) - latents = (1 - init_mask) * init_latents_proper + init_mask * latents + latents = ( + 1 - init_mask + ) * init_latents_proper + init_mask * latents if callback_on_step_end is not None: callback_kwargs = {} @@ -411,12 +445,18 @@ def __call__( latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) mask = callback_outputs.pop("mask", mask) - masked_image_latents = callback_outputs.pop("masked_image_latents", masked_image_latents) + masked_image_latents = callback_outputs.pop( + "masked_image_latents", masked_image_latents + ) # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): progress_bar.update() if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) @@ -425,15 +465,27 @@ def __call__( if not output_type == "latent": condition_kwargs = {} if isinstance(self.vae, AsymmetricAutoencoderKL): - init_image = init_image.to(device=device, dtype=masked_image_latents.dtype) + init_image = init_image.to( + device=device, dtype=masked_image_latents.dtype + ) init_image_condition = init_image.clone() init_image = self._encode_vae_image(init_image, generator=generator) - mask_condition = mask_condition.to(device=device, dtype=masked_image_latents.dtype) - condition_kwargs = {"image": init_image_condition, "mask": mask_condition} + mask_condition = mask_condition.to( + device=device, dtype=masked_image_latents.dtype + ) + condition_kwargs = { + "image": init_image_condition, + "mask": mask_condition, + } image = self.vae.decode( - latents / self.vae.config.scaling_factor, return_dict=False, generator=generator, **condition_kwargs + latents / self.vae.config.scaling_factor, + return_dict=False, + generator=generator, + **condition_kwargs, )[0] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + image, has_nsfw_concept = self.run_safety_checker( + image, device, prompt_embeds.dtype + ) else: image = latents has_nsfw_concept = None @@ -443,10 +495,17 @@ def __call__( else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + image = self.image_processor.postprocess( + image, output_type=output_type, do_denormalize=do_denormalize + ) if padding_mask_crop is not None: - image = [self.image_processor.apply_overlay(mask_image, original_image, i, crops_coords) for i in image] + image = [ + self.image_processor.apply_overlay( + mask_image, original_image, i, crops_coords + ) + for i in image + ] # Offload all models self.maybe_free_model_hooks() @@ -454,10 +513,21 @@ def __call__( if not return_dict: return (image, has_nsfw_concept) - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) + return StableDiffusionPipelineOutput( + images=image, nsfw_content_detected=has_nsfw_concept + ) def prepare_mask_latents( - self, mask, masked_image, batch_size, height, width, dtype, device, generator, cloth_classifier_free_guidance + self, + mask, + masked_image, + batch_size, + height, + width, + dtype, + device, + generator, + cloth_classifier_free_guidance, ): # resize the mask to latents shape as we concatenate the mask to the latents # we do that before converting to dtype to avoid breaking in case we're using cpu_offload @@ -472,7 +542,9 @@ def prepare_mask_latents( if masked_image.shape[1] == 4: masked_image_latents = masked_image else: - masked_image_latents = self._encode_vae_image(masked_image, generator=generator) + masked_image_latents = self._encode_vae_image( + masked_image, generator=generator + ) # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method if mask.shape[0] < batch_size: @@ -490,11 +562,15 @@ def prepare_mask_latents( f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed." " Make sure the number of images that you pass is divisible by the total requested batch size." ) - masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + masked_image_latents = masked_image_latents.repeat( + batch_size // masked_image_latents.shape[0], 1, 1, 1 + ) mask = torch.cat([mask] * 2) if cloth_classifier_free_guidance else mask masked_image_latents = ( - torch.cat([masked_image_latents] * 2) if cloth_classifier_free_guidance else masked_image_latents + torch.cat([masked_image_latents] * 2) + if cloth_classifier_free_guidance + else masked_image_latents ) # aligning device to prevent device errors when concating it with the latent model input diff --git a/pipelines/OmsDiffusionPipeline.py b/pipelines/OmsDiffusionPipeline.py index 4c3ebd5..bf5c10c 100644 --- a/pipelines/OmsDiffusionPipeline.py +++ b/pipelines/OmsDiffusionPipeline.py @@ -4,30 +4,30 @@ class OmsDiffusionPipeline(StableDiffusionPipeline): @torch.no_grad() def __call__( - self, - prompt: Union[str, List[str]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - timesteps: List[int] = None, - guidance_scale: float = 5., - cloth_guidance_scale: float = 2.5, - negative_prompt: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - ip_adapter_image: Optional[PipelineImageInput] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - guidance_rescale: float = 0.0, - clip_skip: Optional[int] = None, - callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, - callback_on_step_end_tensor_inputs: List[str] = ["latents"], - **kwargs, + self, + prompt: Union[str, List[str]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + timesteps: List[int] = None, + guidance_scale: float = 5.0, + cloth_guidance_scale: float = 2.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + guidance_rescale: float = 0.0, + clip_skip: Optional[int] = None, + callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, ): r""" The call function to the pipeline for generation. @@ -157,7 +157,9 @@ def __call__( # 3. Encode input prompt lora_scale = ( - self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + self.cross_attention_kwargs.get("scale", None) + if self.cross_attention_kwargs is not None + else None ) prompt_embeds, negative_prompt_embeds = self.encode_prompt( @@ -176,7 +178,9 @@ def __call__( # Here we concatenate the unconditional and text embeddings into a single batch # to avoid doing two forward passes if self.do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, negative_prompt_embeds, prompt_embeds]) + prompt_embeds = torch.cat( + [negative_prompt_embeds, negative_prompt_embeds, prompt_embeds] + ) if ip_adapter_image is not None: image_embeds = self.prepare_ip_adapter_image_embeds( @@ -184,7 +188,9 @@ def __call__( ) # 4. Prepare timesteps - timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps + ) # 5. Prepare latent variables num_channels_latents = self.unet.config.in_channels @@ -203,12 +209,16 @@ def __call__( extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) # 6.1 Add image embeds for IP-Adapter - added_cond_kwargs = {"image_embeds": image_embeds} if ip_adapter_image is not None else None + added_cond_kwargs = ( + {"image_embeds": image_embeds} if ip_adapter_image is not None else None + ) # 6.2 Optionally get Guidance Scale Embedding timestep_cond = None if self.unet.config.time_cond_proj_dim is not None: - guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat( + batch_size * num_images_per_prompt + ) timestep_cond = self.get_guidance_scale_embedding( guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim ).to(device=device, dtype=latents.dtype) @@ -222,8 +232,14 @@ def __call__( continue # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 3) if self.do_classifier_free_guidance else latents - latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + latent_model_input = ( + torch.cat([latents] * 3) + if self.do_classifier_free_guidance + else latents + ) + latent_model_input = self.scheduler.scale_model_input( + latent_model_input, t + ) # predict the noise residual noise_pred = self.unet( @@ -238,7 +254,9 @@ def __call__( # perform guidance if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_cloth, noise_pred_text = noise_pred.chunk(3) + noise_pred_uncond, noise_pred_cloth, noise_pred_text = ( + noise_pred.chunk(3) + ) noise_pred = ( noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_cloth) @@ -247,10 +265,16 @@ def __call__( if self.do_classifier_free_guidance and self.guidance_rescale > 0.0: # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf - noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale) + noise_pred = rescale_noise_cfg( + noise_pred, + noise_pred_text, + guidance_rescale=self.guidance_rescale, + ) # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] + latents = self.scheduler.step( + noise_pred, t, latents, **extra_step_kwargs, return_dict=False + )[0] if callback_on_step_end is not None: callback_kwargs = {} @@ -260,20 +284,28 @@ def __call__( latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + negative_prompt_embeds = callback_outputs.pop( + "negative_prompt_embeds", negative_prompt_embeds + ) # call the callback, if provided - if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + if i == len(timesteps) - 1 or ( + (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 + ): progress_bar.update() if callback is not None and i % callback_steps == 0: step_idx = i // getattr(self.scheduler, "order", 1) callback(step_idx, t, latents) if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ - 0 - ] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + image = self.vae.decode( + latents / self.vae.config.scaling_factor, + return_dict=False, + generator=generator, + )[0] + image, has_nsfw_concept = self.run_safety_checker( + image, device, prompt_embeds.dtype + ) else: image = latents has_nsfw_concept = None @@ -283,7 +315,9 @@ def __call__( else: do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] - image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + image = self.image_processor.postprocess( + image, output_type=output_type, do_denormalize=do_denormalize + ) # Offload all models self.maybe_free_model_hooks() @@ -291,4 +325,6 @@ def __call__( if not return_dict: return (image, has_nsfw_concept) - return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) + return StableDiffusionPipelineOutput( + images=image, nsfw_content_detected=has_nsfw_concept + ) diff --git a/utils/resampler.py b/utils/resampler.py index 2426667..0056e4e 100644 --- a/utils/resampler.py +++ b/utils/resampler.py @@ -69,7 +69,9 @@ def forward(self, x, latents): # attention scale = 1 / math.sqrt(math.sqrt(self.dim_head)) - weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = (q * scale) @ (k * scale).transpose( + -2, -1 + ) # More stable with f16 than dividing afterwards weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) out = weight @ v @@ -94,7 +96,9 @@ def __init__( num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence ): super().__init__() - self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None + self.pos_emb = ( + nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None + ) self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) @@ -135,7 +139,11 @@ def forward(self, x): x = self.proj_in(x) if self.to_latents_from_mean_pooled_seq: - meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool)) + meanpooled_seq = masked_mean( + x, + dim=1, + mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool), + ) meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq) latents = torch.cat((meanpooled_latents, latents), dim=-2) diff --git a/utils/utils.py b/utils/utils.py index 4dc2078..f0ec648 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -15,7 +15,9 @@ def prepare_image(image, height, width): if isinstance(image, torch.Tensor): # Batch single image if image.ndim == 3: - assert image.shape[0] == 3, "Image outside a batch should be of shape (3, H, W)" + assert ( + image.shape[0] == 3 + ), "Image outside a batch should be of shape (3, H, W)" image = image.unsqueeze(0) # Check image is in [-1, 1] @@ -30,7 +32,9 @@ def prepare_image(image, height, width): image = [image] if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): # resize all images w.r.t passed height an width - image = [i.resize((width, height), resample=PIL.Image.LANCZOS) for i in image] + image = [ + i.resize((width, height), resample=PIL.Image.LANCZOS) for i in image + ] image = [np.array(i.convert("RGB"))[None, :] for i in image] image = np.concatenate(image, axis=0) elif isinstance(image, list) and isinstance(image[0], np.ndarray): @@ -49,7 +53,9 @@ def prepare_mask(image, height, width): if isinstance(image, torch.Tensor): # Batch single image if image.ndim == 3: - assert image.shape[0] == 1, "Image outside a batch should be of shape (3, H, W)" + assert ( + image.shape[0] == 1 + ), "Image outside a batch should be of shape (3, H, W)" image = image.unsqueeze(0) image = image.to(dtype=torch.float32) else: @@ -58,14 +64,16 @@ def prepare_mask(image, height, width): image = [image] if isinstance(image, list) and isinstance(image[0], PIL.Image.Image): # resize all images w.r.t passed height an width - image = [i.resize((width, height), resample=PIL.Image.NEAREST) for i in image] + image = [ + i.resize((width, height), resample=PIL.Image.NEAREST) for i in image + ] image = [np.array(i.convert("L"))[..., None] for i in image] image = np.stack(image, axis=0) elif isinstance(image, list) and isinstance(image[0], np.ndarray): image = np.stack([i[..., None] for i in image], axis=0) image = image.transpose(0, 3, 1, 2) - image = torch.from_numpy(image).to(dtype=torch.float32) / 255. + image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0 image[image > 0.5] = 1 image[image <= 0.5] = 0