blanchon commited on
Commit
4ba46c8
·
1 Parent(s): 2a80021
Files changed (1) hide show
  1. app-fast.py +4 -4
app-fast.py CHANGED
@@ -39,7 +39,7 @@ RESOLUTION_OPTIONS: list[str] = [
39
  device = torch.device("cuda")
40
 
41
  quant_config = TransformersBitsAndBytesConfig(
42
- load_in_8bit=True,
43
  )
44
 
45
 
@@ -51,17 +51,17 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
51
  low_cpu_mem_usage=True,
52
  quantization_config=quant_config,
53
  torch_dtype=torch.float16,
54
- ).to(device) # Move model to the correct device after loading
55
 
56
  quant_config = DiffusersBitsAndBytesConfig(
57
- load_in_8bit=True,
58
  )
59
  transformer = HiDreamImageTransformer2DModel.from_pretrained(
60
  MODEL_PATH,
61
  subfolder="transformer",
62
  quantization_config=quant_config,
63
  torch_dtype=torch.float16,
64
- ).to(device)
65
 
66
  scheduler = MODEL_CONFIGS["scheduler"](
67
  num_train_timesteps=1000,
 
39
  device = torch.device("cuda")
40
 
41
  quant_config = TransformersBitsAndBytesConfig(
42
+ load_in_4bit=True,
43
  )
44
 
45
 
 
51
  low_cpu_mem_usage=True,
52
  quantization_config=quant_config,
53
  torch_dtype=torch.float16,
54
+ )
55
 
56
  quant_config = DiffusersBitsAndBytesConfig(
57
+ load_in_4bit=True,
58
  )
59
  transformer = HiDreamImageTransformer2DModel.from_pretrained(
60
  MODEL_PATH,
61
  subfolder="transformer",
62
  quantization_config=quant_config,
63
  torch_dtype=torch.float16,
64
+ )
65
 
66
  scheduler = MODEL_CONFIGS["scheduler"](
67
  num_train_timesteps=1000,