svg server update (83074896) · Commits · 张泽凯 / VAGEN

vagen/env/init.py

+4 −3

Original line number	Diff line number	Diff line
		@@ -2,7 +2,7 @@ from .sokoban import SokobanEnv,SokobanConfig
		from .frozenlake import FrozenLakeEnv,FrozenLakeConfig, FrozenLakeService
		from .navigation import NavigationEnv, NavigationConfig
		from .svgdino import SVGDINOEnv, SVGDINOConfig
		from .svg import SVGEnv, SVGConfig
		from .svg import SVGEnv, SVGConfig, SVGService

		REGISTERED_ENV = {
		"sokoban": {
		@@ -20,10 +20,11 @@ REGISTERED_ENV = {
		},
		"svg": {
		"env_cls": SVGEnv,
		"config_cls": SVGConfig
		"config_cls": SVGConfig,
		"service_cls": SVGService
		},
		"svgdino": {
		"env_cls": SVGDINOEnv,
		"config_cls": SVGDINOConfig
		"config_cls": SVGDINOConfig,
		}
		}
		No newline at end of file

vagen/env/client.py

+6 −5

Original line number	Diff line number	Diff line
		from typing import Dict, List, Tuple, Optional, Any, Union
		import requests
		import time
		from vagen.utils.serial import deserialize_observation
		from vagen.utils.serial import deserialize_observation, deserialize_step_result

		class BatchEnvClient:
		"""
		@@ -54,8 +54,9 @@ class BatchEnvClient:
		response.raise_for_status() # Raise an exception for 4XX/5XX responses
		return response.json()

		except requests.exceptions.RequestException as e:
		raise ConnectionError(f"Failed to communicate with server: {str(e)}")
		except Exception as e:
		print(f"Exception in _make_request: {str(e)}")
		raise

		def check_server_health(self) -> Dict[str, Any]:
		"""
		@@ -146,8 +147,8 @@ class BatchEnvClient:

		# Deserialize observations
		deserialized_results = {}
		for env_id, (observation, reward, done, info) in results.items():
		deserialized_results[env_id] = (deserialize_observation(observation), reward, done, info)
		for env_id, serialized_result in results.items():
		deserialized_results[env_id] = deserialize_step_result(serialized_result)

		return deserialized_results

vagen/env/svg/init.py

+1 −0

Original line number	Diff line number	Diff line
		from .config import SVGConfig
		from .env import SVGEnv
		from .service import SVGService
		No newline at end of file

vagen/env/svg/dino.py

+26 −0

Original line number	Diff line number	Diff line
		@@ -4,9 +4,35 @@ from tqdm import tqdm
		from transformers import AutoModel, AutoImageProcessor
		from PIL import Image
		import torch.nn as nn
		import threading
		import logging

		# @TODO clean codes of this section

		_model_cache = {}
		_model_cache_lock = threading.Lock()

		def get_dino_model(model_size="small", device="cuda"):
		"""
		Get or create a DINO model instance with singleton pattern

		Args:
		model_size: Size of DINO model ('small', 'base', or 'large')
		device: Device to run model on ('cuda' or 'cpu')

		Returns:
		DINOScoreCalculator instance
		"""
		# Use the actual DINOScoreCalculator implementation from your code
		# with added singleton pattern
		cache_key = f"{model_size}_{device}"

		with _model_cache_lock:
		if cache_key not in _model_cache:
		logging.info(f"Creating new DINO model: {model_size} on {device}")
		_model_cache[cache_key] = DINOScoreCalculator(model_size=model_size, device=device)
		return _model_cache[cache_key]

		class AverageMeter(object):
		"""Computes and stores the average and current value"""

vagen/env/svg/env.py

+12 −49

Original line number	Diff line number	Diff line
		@@ -35,6 +35,7 @@ class SVGEnv(BaseEnv):
		self.gt_image = None
		self.gen_svg_code = None
		self.gen_image = None
		self.dino_model = None

		# Initialize random number generator
		self.rng = random.Random()
		@@ -76,7 +77,7 @@ class SVGEnv(BaseEnv):

		return self._render(init_obs=True), {}

		def step(self, action_str: str) -> Tuple[Dict, float, bool, Dict]:
		def step(self, action_str: str, dino_model=None) -> Tuple[Dict, float, bool, Dict]:
		"""Execute a step in the environment"""
		# Parse LLM response
		rst = parse_llm_raw_response(
		@@ -119,16 +120,6 @@ class SVGEnv(BaseEnv):
		# Invalid format - apply penalty
		self.reward += self.config.format_penalty

		# Log failure if analysis mode is enabled
		if hasattr(self, 'failure_logger'):
		failure_info = {
		'img_id': self.img_id,
		'gt_svg_code': self.gt_svg_code,
		'gen_svg_code': action_str,
		'failure_reason': 'invalid_svg'
		}
		self.failure_logger.info(json.dumps(failure_info))

		done = True
		info["metrics"] = metrics
		self.total_reward += self.reward
		@@ -144,59 +135,28 @@ class SVGEnv(BaseEnv):
		# Process the generated SVG code
		_, gen_image = process_and_rasterize_svg(self.gen_svg_code)
		self.gen_image = gen_image

		# Calculate score
		score_config = {
		"model_size": self.config.model_size,
		"dino_only": self.config.dino_only,
		}

		# Add optional weights if set
		if self.config.dino_weight is not None:
		score_config["dino_weight"] = self.config.dino_weight
		if self.config.structural_weight is not None:
		score_config["structural_weight"] = self.config.structural_weight
		if self.config.color_weight is not None:
		score_config["color_weight"] = self.config.color_weight
		if self.config.code_weight is not None:
		score_config["code_weight"] = self.config.code_weight

		score_config = self.config.get_score_config()
		print(f"score_config:{score_config}")
		scores = calculate_total_score(
		gt_im=self.gt_image,
		gen_im=gen_image,
		gt_code=self.gt_svg_code,
		gen_code=self.gen_svg_code,
		score_config=score_config
		score_config=score_config,
		dino_model=dino_model
		)

		# Set metrics and update reward
		self.reward += scores["total_score"]
		info["scores"] = scores

		# SVG generation is considered effective if score is above threshold
		metrics["turn_metrics"]["action_is_effective"] = scores["total_score"] > 0

		# Log success if analysis mode is enabled
		if hasattr(self, 'success_logger'):
		success_info = {
		'img_id': self.img_id,
		'gt_svg_code': self.gt_svg_code,
		'gen_svg_code': self.gen_svg_code,
		'scores': scores
		}
		self.success_logger.info(json.dumps(success_info))

		except Exception as e:
		# Error processing SVG - log failure
		if hasattr(self, 'failure_logger'):
		failure_info = {
		'img_id': self.img_id,
		'gt_svg_code': self.gt_svg_code,
		'gen_svg_code': self.gen_svg_code,
		'failure_reason': str(e)
		}
		self.failure_logger.info(json.dumps(failure_info))

		import traceback
		print(f"Error processing SVG: {e}")
		traceback.print_exc()
		# Reset actions and update metrics
		self.valid_actions = []
		metrics["turn_metrics"]["action_is_valid"] = False
		@@ -300,6 +260,9 @@ class SVGEnv(BaseEnv):
		success_handler.setFormatter(logging.Formatter('%(message)s'))
		self.success_logger.addHandler(success_handler)

		def set_dino_model(self, model):
		self.dino_model = model

		if __name__ == "__main__":
		config = SVGConfig(
		dataset_name="starvector/svg-emoji-simple",

Admin message