ranamhamoud commited on
Commit
9363edf
·
verified ·
1 Parent(s): 44a3700

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -2
app.py CHANGED
@@ -36,7 +36,11 @@ def generate_image(text):
36
  logging.error("Failed to generate image: %s", str(error))
37
  raise gr.Error("An error occurred while generating the image. Please check your API key and try again.")
38
 
39
-
 
 
 
 
40
 
41
  # Constants
42
  MAX_MAX_NEW_TOKENS = 2048
@@ -63,7 +67,12 @@ if torch.cuda.is_available():
63
  bnb_4bit_quant_type="nf4",
64
  bnb_4bit_compute_dtype=torch.bfloat16
65
  )
66
- base_model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", quantization_config=bnb_config)
 
 
 
 
 
67
  model = PeftModel.from_pretrained(base_model, "ranamhamoud/storytellai-2.0")
68
  tokenizer = AutoTokenizer.from_pretrained(model_id)
69
  tokenizer.pad_token = tokenizer.eos_token
 
36
  logging.error("Failed to generate image: %s", str(error))
37
  raise gr.Error("An error occurred while generating the image. Please check your API key and try again.")
38
 
39
+ rope_scaling = {
40
+ 'type': 'linear', # Adjust the type to the appropriate scaling type for your model.
41
+ 'factor': 8.0 # Use the intended scaling factor.
42
+ }
43
+
44
 
45
  # Constants
46
  MAX_MAX_NEW_TOKENS = 2048
 
67
  bnb_4bit_quant_type="nf4",
68
  bnb_4bit_compute_dtype=torch.bfloat16
69
  )
70
+ base_model = AutoModelForCausalLM.from_pretrained(
71
+ model_id,
72
+ device_map="auto",
73
+ quantization_config=bnb_config,
74
+ rope_scaling=rope_scaling # Add this only if your model specifically requires it.
75
+ )
76
  model = PeftModel.from_pretrained(base_model, "ranamhamoud/storytellai-2.0")
77
  tokenizer = AutoTokenizer.from_pretrained(model_id)
78
  tokenizer.pad_token = tokenizer.eos_token