Update src/pipeline.py

src/pipeline.py

@@ -39,83 +39,16 @@ from diffusers.utils.import_utils import is_torch_npu_available
 from diffusers.utils.torch_utils import maybe_allow_in_graph
 from diffusers.models.embeddings import CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, FluxPosEmbed
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
-
+from diffusers import FluxPipeline, FluxTransformer2DModel
 from model import E, D
 import torchvision 
-
 import os
+
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
 torch._dynamo.config.suppress_errors = True
 
-class BasicQuantization:
-    def __init__(self, bits=1):
-        self.bits = bits
-        self.qmin = -(2**(bits-1))
-        self.qmax = 2**(bits-1) - 1
-
-    def quantize_tensor(self, tensor):
-        scale = (tensor.max() - tensor.min()) / (self.qmax - self.qmin)
-        zero_point = self.qmin - torch.round(tensor.min() / scale)
-        qtensor = torch.round(tensor / scale + zero_point)
-        qtensor = torch.clamp(qtensor, self.qmin, self.qmax)
-        return (qtensor - zero_point) * scale, scale, zero_point
-
-class ModelQuantization:
-    def __init__(self, model, bits=7):
-        self.model = model
-        self.quant = BasicQuantization(bits)
-
-    def quantize_model(self):
-        for name, module in self.model.named_modules():
-            if isinstance(module, torch.nn.Linear):
-                if hasattr(module, 'weightML'):
-                    quantized_weight, _, _ = self.quant.quantize_tensor(module.weight)
-                    module.weight = torch.nn.Parameter(quantized_weight)
-                if hasattr(module, 'bias') and module.bias is not None:
-                    quantized_bias, _, _ = self.quant.quantize_tensor(module.bias)
-                    module.bias = torch.nn.Parameter(quantized_bias)
-
-
-def inicializar_generador(dispositivo: torch.device, respaldo: torch.Generator = None):
-    if dispositivo.type == "cpu":
-        return torch.Generator(device="cpu").set_state(torch.get_rng_state())
-    elif dispositivo.type == "cuda":
-        return torch.Generator(device=dispositivo).set_state(torch.cuda.get_rng_state())
-    else:
-        if respaldo is None:
-            return inicializar_generador(torch.device("cpu"))
-        else:
-            return respaldo
 
-def calcular_fusion(x: torch.Tensor, info_tome: Dict[str, Any]) -> Tuple[Callable, ...]:
-    alto_original, ancho_original = info_tome["size"]
-    tokens_originales = alto_original * ancho_original
-    submuestreo = int(math.ceil(math.sqrt(tokens_originales // x.shape[1])))
-    argumentos = info_tome["args"]
-    if submuestreo <= argumentos["down"]:
-        ancho = int(math.ceil(ancho_original / submuestreo))
-        alto = int(math.ceil(alto_original / submuestreo))
-        radio = int(x.shape[1] * argumentos["ratio"])
-
-        if argumentos["generator"] is None:
-            argumentos["generator"] = inicializar_generador(x.device)
-        elif argumentos["generator"].device != x.device:
-            argumentos["generator"] = inicializar_generador(x.device, respaldo=argumentos["generator"])
-
-        usar_aleatoriedad = argumentos["rando"]
-        fusion, desfusion = ghanta.emparejamiento_suave_aleatorio_2d(
-            x, ancho, alto, argumentos["sx"], argumentos["sy"], radio,
-            sin_aleatoriedad=not usar_aleatoriedad, generador=argumentos["generator"]
-        )
-    else:
-        fusion, desfusion = (hacer_nada, hacer_nada)
-    fusion_a, desfusion_a = (fusion, desfusion) if argumentos["m1"] else (hacer_nada, hacer_nada)
-    fusion_c, desfusion_c = (fusion, desfusion) if argumentos["m2"] else (hacer_nada, hacer_nada)
-    fusion_m, desfusion_m = (fusion, desfusion) if argumentos["m3"] else (hacer_nada, hacer_nada)
-    return fusion_a, fusion_c, fusion_m, desfusion_a, desfusion_c, desfusion_m
-
-from diffusers import FluxPipeline, FluxTransformer2DModel
 Pipeline = None
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
@@ -127,10 +60,6 @@ ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
 TinyVAE = "madebyollin/taef1"
 TinyVAE_REV = "2d552378e58c9c94201075708d7de4e1163b2689"
 
-os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
-os.environ["TOKENIZERS_PARALLELISM"] = "True"
-torch._dynamo.config.suppress_errors = True
-
 def empty_cache():
     gc.collect()
     torch.cuda.empty_cache()
@@ -172,8 +101,7 @@ def load_pipeline() -> Pipeline:
         pipeline(prompt="insensible, timbale, pothery, electrovital, actinogram, taxis, intracerebellar, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
     return pipeline
 
-
-sample = None
+sample = 1
 @torch.no_grad()
 def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator) -> Image:
     global sample
@@ -181,4 +109,4 @@ def infer(request: TextToImageRequest, pipeline: Pipeline, generator: Generator)
         sample=1
         empty_cache()
     image=pipeline(request.prompt,generator=generator, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256, height=request.height, width=request.width, output_type="pt").images[0]
-    return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))# torchvision.transforms.functional.to_pil_image(image)
+    return torchvision.transforms.functional.to_pil_image(image.to(torch.float32).mul_(2).sub_(1))