minor (6e500ee6) · Commits · 张泽凯 / VAGEN

vagen/env/svg/dreamsim.py

+14 −17

Original line number	Diff line number	Diff line
		@@ -2,12 +2,15 @@ import torch
		from PIL import Image
		import os
		from dreamsim import dreamsim
		import threading
		import logging

		# Create global cache and lock, similar to DINO implementation
		_model_cache = {}
		_model_cache_lock = threading.Lock()
		_model_counter = 0


		def get_dreamsim_model(device=None):
		"""
		Get a singleton instance of DreamSim model, using cache to avoid duplicate loading
		@@ -23,6 +26,10 @@ def get_dreamsim_model(device=None):
		# Choose device based on availability if not specified
		if device is None:
		device = "cuda" if torch.cuda.is_available() else "cpu"
		else:
		# If CUDA was requested but not available, fall back to CPU
		if device == "cuda" and not torch.cuda.is_available():
		device = "cpu"

		# Use device as cache key
		cache_key = f"dreamsim_{device}"
		@@ -35,10 +42,12 @@ def get_dreamsim_model(device=None):
		_model_cache[cache_key] = DreamSimScoreCalculator(device=device)
		return _model_cache[cache_key]


		class DreamSimScoreCalculator:
		"""
		A wrapper class for DreamSim model to calculate similarity scores between images.
		"""

		def __init__(self, pretrained=True, cache_dir="~/.cache", device=None):
		"""
		Initialize DreamSim model.
		@@ -49,14 +58,15 @@ class DreamSimScoreCalculator:
		device: Device to run the model on (defaults to CUDA if available, else CPU)
		"""
		cache_dir = os.path.expanduser(cache_dir)
		if device is None:
		self.device = "cuda" if torch.cuda.is_available() else "cpu"

		# Verify device availability
		if device is None or (device == "cuda" and not torch.cuda.is_available()):
		self.device = "cpu"
		else:
		self.device = device

		# Load model and preprocessor
		self.model, self.preprocess = dreamsim(pretrained=pretrained, cache_dir=cache_dir)
		self.model = self.model.to(self.device)
		self.model, self.preprocess = dreamsim(pretrained=pretrained, cache_dir=cache_dir, device=self.device)

		def calculate_similarity_score(self, gt_im, gen_im):
		"""
		@@ -119,16 +129,3 @@ class DreamSimScoreCalculator:
		scores.append(similarity)

		return scores
		No newline at end of file

		# Helper function to get or initialize DreamSim model
		def get_dreamsim_model(device=None):
		"""
		Get an instance of DreamSim model.

		Args:
		device: Device to run model on

		Returns:
		DreamSimScoreCalculator: Instance of DreamSim calculator
		"""
		return DreamSimScoreCalculator(device=device)
		No newline at end of file

vagen/examples/debug/debug_svg_vision_grpo/env_config.yaml

0 → 100644

+21 −0

Original line number	Diff line number	Diff line
		env1:
		env_name: svg
		env_config:
		split: train
		dino_weight: 1.0
		structural_weight: 1.0
		dreamsim_weight: 1.0
		train_size: 10000
		test_size: 0

		env2:
		env_name: svg
		env_config:
		split: test
		dino_weight: 1.0
		structural_weight: 1.0
		dreamsim_weight: 1.0
		train_size: 0
		test_size: 24

vagen/examples/debug/debug_svg_vision_grpo/run.sh

100644 → 100755

+1 −1

Original line number	Diff line number	Diff line
		@@ -59,7 +59,7 @@ python3 -m vagen.trainer.main_ppo \
		trainer.logger=['console','wandb'] \
		trainer.project_name='vagen_debug' \
		trainer.experiment_name='grpo_mask_loss_svg_vision_debug' \
		trainer.n_gpus_per_node=2 \
		trainer.n_gpus_per_node=4 \
		trainer.nnodes=1 \
		trainer.save_freq=100 \
		trainer.test_freq=20 \

vagen/examples/svg_aico/env_config.yaml

+6 −8

Original line number	Diff line number	Diff line
		@@ -2,9 +2,9 @@ env1:
		env_name: svg
		env_config:
		split: train
		dino_weight: 5.0
		structural_weight: 5.0
		dreamsim_weight: 1.0
		dino_weight: 0.0
		structural_weight: 0.0
		dreamsim_weight: 5.0
		train_size: 1000
		test_size: 0

		@@ -12,10 +12,8 @@ env2:
		env_name: svg
		env_config:
		split: test
		dino_weight: 5.0
		structural_weight: 5.0
		dreamsim_weight: 1.0
		dino_weight: 0.0
		structural_weight: 0.0
		dreamsim_weight: 5.0
		train_size: 0
		test_size: 512
		No newline at end of file

vagen/examples/svg_trico/env_config.yaml

+6 −8

Original line number	Diff line number	Diff line
		@@ -2,9 +2,9 @@ env1:
		env_name: svg
		env_config:
		split: train
		dino_weight: 1.0
		structural_weight: 1.0
		dreamsim_weight: 1.0
		dino_weight: 0.0
		structural_weight: 0.0
		dreamsim_weight: 5.0
		train_size: 1000
		test_size: 0

		@@ -12,10 +12,8 @@ env2:
		env_name: svg
		env_config:
		split: test
		dino_weight: 1.0
		structural_weight: 1.0
		dreamsim_weight: 1.0
		dino_weight: 0.0
		structural_weight: 0.0
		dreamsim_weight: 5.0
		train_size: 0
		test_size: 512
		No newline at end of file

Admin message