File size: 3,686 Bytes
012b86f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import os
import gradio as gr
from diffusers import DiffusionPipeline
import torch
from huggingface_hub import login
# 初始化 Hugging Face 登录
if os.getenv('HF_TOKEN'):
login(token=os.getenv('HF_TOKEN'))
# 模型加载函数
def load_model():
model_id = "IndexTeam/Index-anisora"
print("正在加载模型...")
pipe = DiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
variant="fp16" if torch.cuda.is_available() else None,
use_safetensors=True
)
if torch.cuda.is_available():
print("使用CUDA加速")
pipe = pipe.to("cuda")
pipe.enable_attention_slicing()
if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
pipe.enable_xformers_memory_efficient_attention()
else:
print("警告：未检测到CUDA，使用CPU模式运行")
return pipe
# 加载模型
pipe = load_model()
# 生成图像函数
def generate_image(prompt, negative_prompt, steps=30, guidance_scale=7.5, seed=None):
if seed is not None and seed >= 0:
generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(int(seed))
else:
generator = None
try:
with torch.inference_mode():
print(f"生成图像: {prompt[:50]}...")
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=steps,
guidance_scale=guidance_scale,
generator=generator
).images[0]
return image
except Exception as e:
print(f"生成失败: {str(e)}")
raise gr.Error(f"生成图像时出错: {str(e)}")
# 创建Gradio界面
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# Index-anisora 演示
""")
with gr.Row():
with gr.Column():
prompt = gr.Textbox(
label="提示词",
value="1girl, beautiful anime style, detailed background",
placeholder="描述你想要生成的图像..."
)
negative_prompt = gr.Textbox(
label="负面提示词",
value="low quality, blurry, bad anatomy",
placeholder="描述你不希望在图像中出现的内容..."
)
with gr.Row():
steps = gr.Slider(
label="推理步数",
minimum=10,
maximum=50,
value=30,
step=1
)
guidance_scale = gr.Slider(
label="引导强度",
minimum=1.0,
maximum=20.0,
value=7.5,
step=0.5
)
seed = gr.Number(
label="随机种子",
value=-1,
precision=0
)
submit_btn = gr.Button("生成", variant="primary")
with gr.Column():
output_image = gr.Image(label="生成结果")
# 提交按钮事件
submit_btn.click(
fn=generate_image,
inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
outputs=output_image
)
# 启动应用
if __name__ == "__main__":
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=False
) |