Initial commit with folder contents

src/pipeline.py

@@ -1,8 +1,11 @@
-from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
-from diffusers.image_processor import VaeImageProcessor
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from huggingface_hub.constants import HF_HUB_CACHE
-from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
+import torch
+import math
+from typing import Type, Dict, Any, Tuple, Callable, Optional, Union
+import ghanta 
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import torch
 import torch._dynamo
 import gc
@@ -14,16 +17,11 @@ import time
 from diffusers import DiffusionPipeline
 from torchao.quantization import quantize_, int8_weight_only, fpx_weight_only
 import os
-os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
-
-import torch
-import math
-from typing import Type, Dict, Any, Tuple, Callable, Optional, Union
-import ghanta 
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+from diffusers import FluxPipeline, AutoencoderKL, AutoencoderTiny
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from huggingface_hub.constants import HF_HUB_CACHE
+from transformers import T5EncoderModel, T5TokenizerFast, CLIPTokenizer, CLIPTextModel
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
@@ -41,6 +39,7 @@ 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
+os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 
 class BasicQuantization:
     def __init__(self, bits=1):
@@ -201,7 +200,8 @@ class FluxTransformerBlock(nn.Module):
         self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
         self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
 
-        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        # self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5)
         self.ff_context = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
         self._chunk_size = None
         self._chunk_dim = 0
@@ -433,6 +433,7 @@ class FluxTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOrig
     ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
         if joint_attention_kwargs is not None:
             joint_attention_kwargs = joint_attention_kwargs.copy()
+            # lora_scale = joint_attention_kwargs.pop("scale", 2.0)
             lora_scale = joint_attention_kwargs.pop("scale", 1.0)
         else:
             lora_scale = 1.0
@@ -577,8 +578,6 @@ torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
 torch.backends.cudnn.benchmark = True
 
-# ckpt_id = "black-forest-labs/FLUX.1-schnell"
-# ckpt_revision = "741f7c3ce8b383c54771c7003378a50191e9efe9" 
 ckpt_id = "Charlyki/extra0Laye0"
 ckpt_revision = "2c8c7664a42b9eaf96a979dd2d90910ddb0d1e42" 
 def empty_cache():
@@ -602,11 +601,11 @@ def load_pipeline() -> Pipeline:
     
     path = os.path.join(HF_HUB_CACHE, "models--Charlyki--extra1Laye1/snapshots/e851ea424c0d88d420cf85e5adccd41d406d358f")
     generator = torch.Generator(device=device)
+    # model = torch.compile(model, mode="max-autotune-no-cudagraphs")
+    # model = torch.compile(model,backend="aot_eager")
     model = FluxTransformer2DModel.from_pretrained(path, torch_dtype=dtype, use_safetensors=False, generator= generator).to(memory_format=torch.channels_last)
     torch.backends.cudnn.benchmark = True
     torch.backends.cudnn.deterministic = False
-    # model = torch.compile(model, mode="max-autotune-no-cudagraphs")
-    # model = torch.compile(model,backend="aot_eager")
     vae = torch.compile(vae)
     pipeline = DiffusionPipeline.from_pretrained(
         ckpt_id,
@@ -621,8 +620,7 @@ def load_pipeline() -> Pipeline:
     pipeline.text_encoder_2.requires_grad_(False)
     pipeline.text_encoder.requires_grad_(False)
 
-    # pipeline.enable_sequential_cpu_offload(exclude=["transformer"])
-    
+    # for _ in range(4):
     for _ in range(3):
         pipeline(prompt="blah blah waah waah oneshot oneshot gang gang", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)