flaxdiff 0.2.6.1__tar.gz → 0.2.7__tar.gz

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.2.6.1
+Version: 0.2.7
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/flaxdiff/data/sources/images.py

@@ -113,12 +113,14 @@ class ImageTFDSSource(DataSource):
 class ImageTFDSAugmenter(DataAugmenter):
     """Augmenter for TFDS image datasets."""
     
-    def __init__(self, label_path: str = "/home/mrwhite0racle/tensorflow_datasets/oxford_flowers102/2.1.1/label.labels.txt"):
+    def __init__(self, label_path: str = None):
         """Initialize a TFDS image augmenter.
         
         Args:
             label_path: Path to the labels file for datasets like Oxford Flowers.
         """
+        if label_path is None:
+            label_path = os.path.join(os.path.expanduser("~"), "tensorflow_datasets/oxford_flowers102/2.1.1/label.labels.txt")
         self.label_path = label_path
     
     def create_transform(self, image_scale: int = 256, method: Any = None) -> Callable[[], pygrain.MapTransform]:

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/flaxdiff/inference/pipeline.py

@@ -25,6 +25,7 @@ from flaxdiff.inference.utils import parse_config, load_from_wandb_run, load_fro
 @dataclass
 class InferencePipeline:
     """Inference pipeline for a general model."""
+    name: str = None
     model: nn.Module = None
     state: SimpleTrainState = None
     best_state: SimpleTrainState = None
@@ -44,6 +45,7 @@ class DiffusionInferencePipeline(InferencePipeline):
     This pipeline handles loading models from wandb and generating samples using the
     DiffusionSampler from FlaxDiff.
     """
+    artifact: Any = None
     state: TrainState = None
     best_state: TrainState = None
     rngstate: Optional[RandomMarkovState] = None
@@ -51,7 +53,6 @@ class DiffusionInferencePipeline(InferencePipeline):
     model_output_transform: DiffusionPredictionTransform = None
     autoencoder: AutoEncoder = None
     input_config: DiffusionInputConfig = None
-    wandb_run = None
     samplers: Dict[Type[DiffusionSampler], Dict[float, DiffusionSampler]] = field(default_factory=dict)
     config: Dict[str, Any] = field(default_factory=dict)
     
@@ -76,7 +77,7 @@ class DiffusionInferencePipeline(InferencePipeline):
         Returns:
             DiffusionInferencePipeline instance
         """
-        states, config, run = load_from_wandb_run(
+        states, config, run, artifact = load_from_wandb_run(
             wandb_run,
             project=project,
             entity=entity,
@@ -95,6 +96,7 @@ class DiffusionInferencePipeline(InferencePipeline):
             best_state=best_state,
             rngstate=RandomMarkovState(jax.random.PRNGKey(42)),
             run=run,
+            artifact=artifact,
         )
         return pipeline
     
@@ -119,7 +121,7 @@ class DiffusionInferencePipeline(InferencePipeline):
         Returns:
             DiffusionInferencePipeline instance
         """
-        states, config, run = load_from_wandb_registry(
+        states, config, run, artifact = load_from_wandb_registry(
             modelname=modelname,
             project=project,
             entity=entity,
@@ -140,6 +142,7 @@ class DiffusionInferencePipeline(InferencePipeline):
             best_state=best_state,
             rngstate=RandomMarkovState(jax.random.PRNGKey(42)),
             run=run,
+            artifact=artifact,
         )
         return pipeline
             
@@ -151,11 +154,14 @@ class DiffusionInferencePipeline(InferencePipeline):
         best_state: Optional[Dict[str, Any]] = None,
         rngstate: Optional[RandomMarkovState] = None,
         run=None,
+        artifact=None,
     ):
         if rngstate is None:
             rngstate = RandomMarkovState(jax.random.PRNGKey(42))
         # Build and return pipeline
         return cls(
+            name=run.name if run else None,
+            artifact=artifact,
             model=config['model'],
             state=state,
             best_state=best_state,
@@ -165,7 +171,6 @@ class DiffusionInferencePipeline(InferencePipeline):
             autoencoder=config['autoencoder'],
             input_config=config['input_config'],
             config=config,
-            wandb_run=run,
         )
     
     def get_sampler(

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/flaxdiff/inference/utils.py

@@ -292,7 +292,7 @@ def load_from_wandb_run(
         config = run.config
     except Exception as e:
         print(f"Warning: Failed to load model from wandb: {e}")
-    return states, config, run
+    return states, config, run, artifact
 
 def load_from_wandb_registry(
     modelname: str,
@@ -318,4 +318,4 @@ def load_from_wandb_registry(
         config = run.config
     except Exception as e:
         print(f"Warning: Failed to load model from wandb: {e}")
-    return states, config, run
+    return states, config, run, artifact

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.2.6.1
+Version: 0.2.7
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.2.6.1 → flaxdiff-0.2.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.2.6.1"
+version = "0.2.7"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [