up

app-fast.py

@@ -39,7 +39,7 @@ RESOLUTION_OPTIONS: list[str] = [
 device = torch.device("cuda")
 
 quant_config = TransformersBitsAndBytesConfig(
-    load_in_8bit=True,
+    load_in_4bit=True,
 )
 
 
@@ -51,17 +51,17 @@ text_encoder = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True,
     quantization_config=quant_config,
     torch_dtype=torch.float16,
-).to(device)  # Move model to the correct device after loading
+)
 
 quant_config = DiffusersBitsAndBytesConfig(
-    load_in_8bit=True,
+    load_in_4bit=True,
 )
 transformer = HiDreamImageTransformer2DModel.from_pretrained(
     MODEL_PATH,
     subfolder="transformer",
     quantization_config=quant_config,
     torch_dtype=torch.float16,
-).to(device)
+)
 
 scheduler = MODEL_CONFIGS["scheduler"](
     num_train_timesteps=1000,