Upload folder using huggingface_hub

.gitignore

@@ -5,4 +5,4 @@
 **/.venv
 .venv
 .git
-
+*.swp

src/pipeline.py

@@ -10,14 +10,12 @@ from torch import Generator
 from torchao.quantization import quantize_, int8_weight_only
 from time import perf_counter
 
-FLUX_CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
+
 HOME = os.environ["HOME"]
-# REPO_DIR = ".cache/huggingface/hub/models--jokerbit--flux-schnell-int8/snapshots/9510dd83d6d44ab375b5e8facec10afa81be2a8f"
-QUANTIZED_MODEL = ["transformer", "text_encoder", "text_encoder_2", "vae"]
-# QUANT_CKPT = {"transformer": os.path.join(HOME, REPO_DIR, "flux_schnell_transformer_int8wo.pt"),
-#               "text_encoder": os.path.join(HOME, REPO_DIR, "flux_schnell_text_encoder_int8wo.pt"), 
-#               "text_encoder_2": os.path.join(HOME, REPO_DIR, "flux_schnell_text_encoder_2_int8wo.pt"), 
-#               "vae": os.path.join(HOME, REPO_DIR, "flux_schnell_vae_int8wo.pt")}
+FLUX_CHECKPOINT = os.path.join(HOME, 
+                               ".cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/snapshots/741f7c3ce8b383c54771c7003378a50191e9efe9/")
+QUANTIZED_MODEL = ["transformer", "text_encoder_2", "text_encoder", "vae"]
+
 
 QUANT_CONFIG = int8_weight_only()
 DTYPE = torch.bfloat16
@@ -25,7 +23,7 @@ NUM_STEPS = 4
 
 def get_transformer(quantize: bool = True, quant_config = int8_weight_only(), quant_ckpt: str = None):
     if quant_ckpt is not None:
-        config = FluxTransformer2DModel.load_config(FLUX_CHECKPOINT, subfolder="transformer")
+        config = FluxTransformer2DModel.load_config(FLUX_CHECKPOINT, subfolder="transformer", local_files_only=True)
         model = FluxTransformer2DModel.from_config(config).to(DTYPE)
         state_dict = torch.load(quant_ckpt, map_location="cpu")
         model.load_state_dict(state_dict, assign=True)
@@ -33,7 +31,7 @@ def get_transformer(quantize: bool = True, quant_config = int8_weight_only(), qu
         return model
     
     model = FluxTransformer2DModel.from_pretrained(
-        FLUX_CHECKPOINT, subfolder="transformer", torch_dtype=DTYPE,
+        FLUX_CHECKPOINT, subfolder="transformer", torch_dtype=DTYPE, local_files_only=True
         )
     if quantize:
         quantize_(model, quant_config)
@@ -42,7 +40,7 @@ def get_transformer(quantize: bool = True, quant_config = int8_weight_only(), qu
 
 def get_text_encoder(quantize: bool = True, quant_config = int8_weight_only(), quant_ckpt: str = None):
     if quant_ckpt is not None:
-        config = CLIPTextConfig.from_pretrained(FLUX_CHECKPOINT, subfolder="text_encoder")
+        config = CLIPTextConfig.from_pretrained(FLUX_CHECKPOINT, subfolder="text_encoder", local_files_only=True)
         model = CLIPTextModel(config).to(DTYPE)
         state_dict = torch.load(quant_ckpt, map_location="cpu")
         model.load_state_dict(state_dict, assign=True)
@@ -50,7 +48,7 @@ def get_text_encoder(quantize: bool = True, quant_config = int8_weight_only(), q
         return model
     
     model = CLIPTextModel.from_pretrained(
-        FLUX_CHECKPOINT, subfolder="text_encoder", torch_dtype=DTYPE
+        FLUX_CHECKPOINT, subfolder="text_encoder", torch_dtype=DTYPE, local_files_only=True
         )
     if quantize:
         quantize_(model, quant_config)
@@ -59,7 +57,7 @@ def get_text_encoder(quantize: bool = True, quant_config = int8_weight_only(), q
 
 def get_text_encoder_2(quantize: bool = True, quant_config = int8_weight_only(), quant_ckpt: str = None):
     if quant_ckpt is not None:
-        config = T5Config.from_pretrained(FLUX_CHECKPOINT, subfolder="text_encoder_2")
+        config = T5Config.from_pretrained(FLUX_CHECKPOINT, subfolder="text_encoder_2", local_files_only=True)
         model = T5EncoderModel(config).to(DTYPE)
         state_dict = torch.load(quant_ckpt, map_location="cpu")
         print(f"Loaded {quant_ckpt}")
@@ -67,7 +65,7 @@ def get_text_encoder_2(quantize: bool = True, quant_config = int8_weight_only(),
         return model
     
     model = T5EncoderModel.from_pretrained(
-        FLUX_CHECKPOINT, subfolder="text_encoder_2", torch_dtype=DTYPE
+        FLUX_CHECKPOINT, subfolder="text_encoder_2", torch_dtype=DTYPE, local_files_only=True
         )
     if quantize:
         quantize_(model, quant_config)
@@ -76,14 +74,14 @@ def get_text_encoder_2(quantize: bool = True, quant_config = int8_weight_only(),
 
 def get_vae(quantize: bool = True, quant_config = int8_weight_only(), quant_ckpt: str = None):
     if quant_ckpt is not None:
-        config = AutoencoderKL.load_config(FLUX_CHECKPOINT, subfolder="vae")
+        config = AutoencoderKL.load_config(FLUX_CHECKPOINT, subfolder="vae", local_files_only=True)
         model = AutoencoderKL.from_config(config).to(DTYPE)
         state_dict = torch.load(quant_ckpt, map_location="cpu")
         model.load_state_dict(state_dict, assign=True)
         print(f"Loaded {quant_ckpt}")
         return model
     model = AutoencoderKL.from_pretrained(
-        FLUX_CHECKPOINT, subfolder="vae", torch_dtype=DTYPE
+        FLUX_CHECKPOINT, subfolder="vae", torch_dtype=DTYPE, local_files_only=True
         )
     if quantize:
         quantize_(model, quant_config)
@@ -119,7 +117,7 @@ def infer(request: TextToImageRequest, _pipeline: FluxPipeline) -> Image:
     if request.seed is None:
         generator = None
     else:
-        generator = Generator(request.seed).device(_pipeline.device)
+        generator = Generator(device=_pipeline.device).manual_seed(request.seed)
 
     empty_cache()
     image = _pipeline(prompt=request.prompt,