# SAAH_WM

**Repository Path**: li_Y_q/SAAH_WM

## Basic Information

- **Project Name**: SAAH_WM
- **Description**: No description available
- **Primary Language**: Python
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2025-07-16
- **Last Updated**: 2025-08-04

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# SAAH-WM: 语义感知与自适应分层水印系统

## 项目概述

SAAH-WM (Semantic-Aware Adaptive Hierarchical Watermarking) 是一个先进的AI生成图像水印系统，旨在为Stable Diffusion等生成模型提供强鲁棒性的水印保护。

## 当前状态：第一步 - 环境与基线验证

本阶段实现了基础的模型加载和图像生成验证功能，为后续水印功能开发奠定基础。

## 项目结构

```
SAAH-WM/
├── models/              # 模型加载模块
│   └── loader.py        # SD2和CLIP模型加载器
├── utils/               # 工具函数
│   └── data_loader.py   # 数据集加载器
├── outputs/             # 生成图像输出目录
├── docs/                # 项目文档
├── baseline_code/       # 参考的开源水印代码
├── generate_baseline.py # 基线验证脚本
├── requirements.txt     # 项目依赖
└── README.md           # 项目说明
```

## 环境要求

### 硬件要求
- GPU: 推荐NVIDIA GPU (8GB+ VRAM)
- RAM: 16GB+
- 存储: 20GB+ 可用空间

### 软件要求
- Python 3.8+
- CUDA 11.8+ (如使用GPU)
- PyTorch 2.0+

## 安装说明

1. **克隆项目**
```bash
git clone <repository-url>
cd SAAH-WM
```

2. **安装依赖**
```bash
pip install -r requirements.txt
```

3. **准备模型文件**
确保以下模型已下载到本地缓存路径 `/media/wang003/liyongqing/SAAH-WM/cache`：
- `models--stabilityai--stable-diffusion-2-base`
- `models--laion--CLIP-ViT-g-14-laion2B-s12B-b42K`

## 使用方法

### 基线验证

运行基线验证脚本，测试环境配置和模型加载：

```bash
python generate_baseline.py
```

**预期输出：**
- 控制台显示详细的加载和生成日志
- `outputs/` 目录生成512x512的PNG图像
- 生成 `generate_baseline.log` 日志文件

### 技术配置

- **扩散模型**: Stable Diffusion 2 base
- **调度器**: DDIM (50步采样)
- **引导尺度**: 7.5
- **图像尺寸**: 512×512
- **数据集**: Gustavosta/stable-diffusion-prompts

## 核心功能

### 1. 模型加载器 (`models/loader.py`)

- `load_stable_diffusion_model()`: 加载SD2模型并配置DDIM调度器
- `load_clip_model()`: 加载CLIP文本编码器
- `check_models_availability()`: 验证本地模型文件

### 2. 数据加载器 (`utils/data_loader.py`)

- `PromptDataLoader`: 提示词数据集管理类
- `get_random_prompt()`: 随机获取提示词
- 支持批量提示词获取和索引访问

### 3. 基线生成器 (`generate_baseline.py`)

- 端到端图像生成验证
- 详细的日志记录
- 自动输出目录管理
- GPU内存优化

## 日志系统

项目使用Python logging模块提供详细的运行日志：

- **控制台输出**: 实时显示关键信息
- **文件日志**: `generate_baseline.log` 保存完整日志
- **日志级别**: INFO级别，包含模型加载、生成进度等

## 故障排除

### 常见问题

1. **模型路径错误**
   - 检查 `/media/wang003/liyongqing/SAAH-WM/cache` 路径是否存在
   - 确认模型目录名称正确

2. **内存不足**
   - 使用较小的batch size
   - 启用attention slicing (已默认开启)
   - 考虑使用CPU模式

3. **网络连接问题**
   - 确保使用 `local_files_only=True` 参数
   - 检查数据集是否需要首次下载

## 开发计划

- [x] **第一步**: 环境与基线验证
- [ ] **第二步**: 语义哈希生成
- [ ] **第三步**: 自适应水印嵌入
- [ ] **第四步**: 分层检测机制
- [ ] **第五步**: 鲁棒性测试

## 贡献指南

1. Fork 项目
2. 创建功能分支
3. 提交更改
4. 发起 Pull Request

## 许可证

本项目采用 MIT 许可证 - 详见 [LICENSE](LICENSE) 文件。

## 联系方式

如有问题或建议，请通过以下方式联系：
- 项目Issues: <repository-issues-url>
- 邮箱: <contact-email>

---

**注意**: 当前版本仅包含基线验证功能，不包含水印逻辑。水印功能将在后续版本中实现。