Update app.py

app.py

@@ -4,55 +4,27 @@ import logging
 import torch
 from PIL import Image
 import spaces
-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())
+cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+os.environ["TRANSFORMERS_CACHE"] = cache_path
+os.environ["HF_HUB_CACHE"] = cache_path
+os.environ["HF_HOME"] = cache_path
 
-    # 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()
+torch.backends.cuda.matmul.allow_tf32 = True
 
-# Print model size
-print(compute_module_sizes(model)[""] / 1024 / 1204)
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
+pipe.to(device="cuda", dtype=torch.bfloat16)
 
-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