Update app.py

app.py

@@ -1,154 +1,134 @@
+import spaces # 必须在最顶部导入
 import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo" # Replace to the model you would like to use
-
-if torch.cuda.is_available():
- torch_dtype = torch.float16
-else:
- torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
- prompt,
- negative_prompt,
- seed,
- randomize_seed,
- width,
- height,
- guidance_scale,
- num_inference_steps,
- progress=gr.Progress(track_tqdm=True),
-):
- if randomize_seed:
- seed = random.randint(0, MAX_SEED)
-
- generator = torch.Generator().manual_seed(seed)
-
- image = pipe(
- prompt=prompt,
- negative_prompt=negative_prompt,
- guidance_scale=guidance_scale,
- num_inference_steps=num_inference_steps,
- width=width,
- height=height,
- generator=generator,
- ).images[0]
-
- return image, seed
-
-
-examples = [
- "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
- "An astronaut riding a green horse",
- "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
- margin: 0 auto;
- max-width: 640px;
-}
-"""
+import os
 
-with gr.Blocks(css=css) as demo:
- with gr.Column(elem_id="col-container"):
- gr.Markdown(" # Text-to-Image Gradio Template")
-
- with gr.Row():
- prompt = gr.Text(
- label="Prompt",
- show_label=False,
- max_lines=1,
- placeholder="Enter your prompt",
- container=False,
- )
+# 获取 Hugging Face 访问令牌
+hf_token = os.getenv("HF_API_TOKEN")
 
- run_button = gr.Button("Run", scale=0, variant="primary")
+# 定义基础模型名称
+base_model_name = "larry1129/WooWoof_AI_Vision_merged_16bit"
 
- result = gr.Image(label="Result", show_label=False)
+# 定义 adapter 模型名称
+adapter_model_name = "larry1129/WooWoof_AI_Vision_merged_16bit"
 
- with gr.Accordion("Advanced Settings", open=False):
- negative_prompt = gr.Text(
- label="Negative prompt",
- max_lines=1,
- placeholder="Enter a negative prompt",
- visible=False,
- )
+# 定义全局变量用于缓存模型和分词器
+model = None
+tokenizer = None
 
- seed = gr.Slider(
- label="Seed",
- minimum=0,
- maximum=MAX_SEED,
- step=1,
- value=0,
+# 定义提示生成函数
+def generate_prompt(instruction, input_text=""):
+ if input_text:
+ prompt = f"""### Instruction:
+{instruction}
+### Input:
+{input_text}
+### Response:
+"""
+ else:
+ prompt = f"""### Instruction:
+{instruction}
+### Response:
+"""
+ return prompt
+
+# 定义生成响应的函数，并使用 @spaces.GPU 装饰
+@spaces.GPU(duration=40) # 建议将 duration 增加到 120
+def generate_response(instruction, input_text):
+ global model, tokenizer
+
+ if model is None:
+ print("开始加载模型...")
+ # 检查 bitsandbytes 是否已安装
+ import importlib.util
+ if importlib.util.find_spec("bitsandbytes") is None:
+ import subprocess
+ subprocess.call(["pip", "install", "--upgrade", "bitsandbytes"])
+
+ try:
+ # 在函数内部导入需要 GPU 的库
+ import torch
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+ from peft import PeftModel
+
+ # 创建量化配置
+ bnb_config = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_use_double_quant=True,
+ bnb_4bit_quant_type="nf4",
+ bnb_4bit_compute_dtype=torch.float16
  )
 
- randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
- with gr.Row():
- width = gr.Slider(
- label="Width",
- minimum=256,
- maximum=MAX_IMAGE_SIZE,
- step=32,
- value=1024, # Replace with defaults that work for your model
- )
-
- height = gr.Slider(
- label="Height",
- minimum=256,
- maximum=MAX_IMAGE_SIZE,
- step=32,
- value=1024, # Replace with defaults that work for your model
- )
-
- with gr.Row():
- guidance_scale = gr.Slider(
- label="Guidance scale",
- minimum=0.0,
- maximum=10.0,
- step=0.1,
- value=0.0, # Replace with defaults that work for your model
- )
-
- num_inference_steps = gr.Slider(
- label="Number of inference steps",
- minimum=1,
- maximum=50,
- step=1,
- value=2, # Replace with defaults that work for your model
- )
-
- gr.Examples(examples=examples, inputs=[prompt])
- gr.on(
- triggers=[run_button.click, prompt.submit],
- fn=infer,
- inputs=[
- prompt,
- negative_prompt,
- seed,
- randomize_seed,
- width,
- height,
- guidance_scale,
- num_inference_steps,
- ],
- outputs=[result, seed],
- )
-
-if __name__ == "__main__":
- demo.launch()
+ # 加载分词器
+ tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
+ print("分词器加载成功。")
+
+ # 加载基础模型
+ base_model = AutoModelForCausalLM.from_pretrained(
+ base_model_name,
+ quantization_config=bnb_config,
+ device_map="auto",
+ use_auth_token=hf_token,
+ trust_remote_code=True
+ )
+ print("基础模型加载成功。")
+
+ # 加载适配器模型
+ model = PeftModel.from_pretrained(
+ base_model,
+ adapter_model_name,
+ torch_dtype=torch.float16,
+ use_auth_token=hf_token
+ )
+ print("适配器模型加载成功。")
+
+ # 设置 pad_token
+ tokenizer.pad_token = tokenizer.eos_token
+ model.config.pad_token_id = tokenizer.pad_token_id
+
+ # 切换到评估模式
+ model.eval()
+ print("模型已切换到评估模式。")
+ except Exception as e:
+ print("加载模型时出错：", e)
+ raise e
+ else:
+ # 在函数内部导入需要的库
+ import torch
+
+ # 检查 model 和 tokenizer 是否已正确加载
+ if model is None or tokenizer is None:
+ print("模型或分词器未正确加载。")
+ raise ValueError("模型或分词器未正确加载。")
+
+ # 生成提示
+ prompt = generate_prompt(instruction, input_text)
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+ with torch.no_grad():
+ outputs = model.generate(
+ input_ids=inputs["input_ids"],
+ attention_mask=inputs.get("attention_mask"),
+ max_new_tokens=128,
+ temperature=0.7,
+ top_p=0.95,
+ do_sample=True,
+ )
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+ response = response.split("### Response:")[-1].strip()
+ return response
+
+# 创建 Gradio 接口
+iface = gr.Interface(
+ fn=generate_response,
+ inputs=[
+ gr.Textbox(lines=2, placeholder="Instruction", label="Instruction"),
+ ],
+ outputs="text",
+ title="WooWoof AI",
+ description="Based on LLAMA 3.1 for pet related",
+ allow_flagging="never"
+)
+
+# 启动 Gradio 接口
+iface.launch(share=True)