Update opensora/serve/gradio_web_server.py

opensora/serve/gradio_web_server.py

@@ -72,8 +72,23 @@ if __name__ == '__main__':
     vae.latent_size = latent_size
     transformer_model.force_images = args.force_images
     tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name)
-    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_name,
-                                                  torch_dtype=torch.float16).to(device)
+
+    load_8bit, load_4bit = True, False
+    kwargs = {"device_map": "auto"}
+    if load_8bit:
+        kwargs['load_in_8bit'] = True
+    elif load_4bit:
+        from transformers import BitsAndBytesConfig
+        kwargs['load_in_4bit'] = True
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+    else:
+        kwargs['torch_dtype'] = torch.float16
+    text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir="cache_dir", **kwargs)
 
     # set eval mode
     transformer_model.eval()