Remove unnecessary import and minor style change

This commit is contained in:
Xrvk 2024-08-18 01:59:30 +03:00 committed by GitHub
parent 2e6ae1240c
commit 90af574794
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -15,7 +15,6 @@ from .layers import (
) )
from einops import rearrange, repeat from einops import rearrange, repeat
import numpy as np
import comfy.ldm.common_dit import comfy.ldm.common_dit
@dataclass @dataclass
@ -33,6 +32,7 @@ class FluxParams:
qkv_bias: bool qkv_bias: bool
guidance_embed: bool guidance_embed: bool
class Flux(nn.Module): class Flux(nn.Module):
""" """
Transformer model for flow matching on sequences. Transformer model for flow matching on sequences.
@ -106,9 +106,9 @@ class Flux(nn.Module):
if self.params.guidance_embed: if self.params.guidance_embed:
if guidance is None: if guidance is None:
raise ValueError("Didn't get guidance strength for guidance distilled model.") raise ValueError("Didn't get guidance strength for guidance distilled model.")
vec.add_(self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))) vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
vec.add_(self.vector_in(y)) vec = vec + self.vector_in(y)
txt = self.txt_in(txt) txt = self.txt_in(txt)
ids = torch.cat((txt_ids, img_ids), dim=1) ids = torch.cat((txt_ids, img_ids), dim=1)