Merge branch 'main' of https://huggingface.co/spaces/Pipatpong/VCM_Demo

app.py

@@ -3,12 +3,15 @@
 import gradio as gr
 import re
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 checkpoint = "Pipatpong/vcm_santa"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "CPU"
+
+quantization_config = BitsAndBytesConfig(load_in_8bit_fp32_cpu_offload=True)
+
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, low_cpu_mem_usage=True, load_in_8bit=True, device_map="auto", quantization_config=quantization_config)
 
 def generate(text, max_length, num_return_sequences=1):
     inputs = tokenizer.encode(text, padding=False, add_special_tokens=False, return_tensors="pt")