Update app.py

app.py

@@ -8,15 +8,57 @@ from diffusers import DiffusionPipeline
 import copy
 import random
 import time
+from huggingface_hub import hf_hub_download
+from accelerate.utils import set_module_tensor_to_device, compute_module_sizes
+from accelerate import init_empty_weights
+from convert_nf4_flux import replace_with_bnb_linear, create_quantized_param, check_quantized_param
+from diffusers import FluxTransformer2DModel, FluxPipeline
+import safetensors.torch
+import gc
+import torch
+
+# Set dtype and check for float8 support
+dtype = torch.bfloat16
+is_torch_e4m3fn_available = hasattr(torch, "float8_e4m3fn")
+
+ckpt_path = hf_hub_download("ABDALLALSWAITI/Maxwell", filename="diffusion_pytorch_model.safetensors")
+original_state_dict = safetensors.torch.load_file(ckpt_path)
+
+with init_empty_weights():
+    config = FluxTransformer2DModel.load_config("ABDALLALSWAITI/Maxwell")
+    model = FluxTransformer2DModel.from_config(config).to(dtype)
+    expected_state_dict_keys = list(model.state_dict().keys())
+
+    # Load the state dict into the quantized model
+for param_name, param in original_state_dict.items():
+    if param_name not in expected_state_dict_keys:
+        continue
+    
+    is_param_float8_e4m3fn = is_torch_e4m3fn_available and param.dtype == torch.float8_e4m3fn
+    if torch.is_floating_point(param) and not is_param_float8_e4m3fn:
+        param = param.to(dtype)
+    
+    if not check_quantized_param(model, param_name):
+        set_module_tensor_to_device(model, param_name, device=0, value=param)
+    else:
+        create_quantized_param(
+            model, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True
+        )
+
+# Clean up
+del original_state_dict
+gc.collect()
+
+# Print model size
+print(compute_module_sizes(model)[""] / 1024 / 1204)
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/flux.1-dev", transformer=model, torch_dtype=dtype)
+pipe.enable_model_cpu_offload()
 
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
-
-# Initialize the base model
-base_model = "sayakpaul/FLUX.1-merged"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-
+    
 MAX_SEED = 2**32-1
 
 class calculateDuration: