flaxdiff 0.2.6.1__py3-none-any.whl → 0.2.7__py3-none-any.whl

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/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/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.dist-info → flaxdiff-0.2.7.dist-info}/METADATA

@@ -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.dist-info → flaxdiff-0.2.7.dist-info}/RECORD

@@ -9,13 +9,13 @@ flaxdiff/data/sources/audio_utils.py,sha256=X27gG1yQt_abVOYgMtruYmZD7-8_uQCRhhTS
 flaxdiff/data/sources/av_example.py,sha256=RIcbVKqckFqbfnV65NQotzIBxjdDuM67kD1nY8fqw5Q,3826
 flaxdiff/data/sources/av_utils.py,sha256=LCr9MJNurOaoxY-sjzkLqJS_MlX0x3gRSlKAVIglAU0,24045
 flaxdiff/data/sources/base.py,sha256=uhF0odJSYRy0SLw1xnI9Q_q_xiVht2DmEYcX1j9AWT4,4246
-flaxdiff/data/sources/images.py,sha256=PSAJaECkXGeO0HWBveuW29AyNxAEBIb2wkNeyZNVJVE,11716
+flaxdiff/data/sources/images.py,sha256=RFLtKW1xzw6ZPVXtCMmnTg1MPb8dc7rP77rZWbK7qpo,11796
 flaxdiff/data/sources/utils.py,sha256=kFzM4_kPoThbAu54ulABmEDAR33tR50NgzXIpC0Dzjk,7316
 flaxdiff/data/sources/videos.py,sha256=CVpOH6A4P2D8iv3gZIhd2GB5ATUD8Vsm_wVYbbugWD4,9359
 flaxdiff/data/sources/voxceleb2.py,sha256=BoKfat_hsw6ObDyyaiQmPbBzuFiqgCGlgAZmf-t5Iz8,18621
 flaxdiff/inference/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-flaxdiff/inference/pipeline.py,sha256=2eQEeQYuoWL6_d_AdTmAkvtz_poDmwPr9fihOnCt43Y,9045
-flaxdiff/inference/utils.py,sha256=0QEkPjFdqTmMsRaSUHgI9GV8gC2uCo1z40PvsnNYWaw,12303
+flaxdiff/inference/pipeline.py,sha256=8S30FAlXEjvrDd87H-qdD6biySQZ3cJUflU8gdmPxig,9223
+flaxdiff/inference/utils.py,sha256=MVnWl0LnC-1ILk0SsLd1YFu6igaQFR7mGhzo0jE797E,12323
 flaxdiff/inputs/__init__.py,sha256=ybPjQsFAf5sqRVZG1sRiOl99EnwpI-NQ8HE3y7UbXmU,7197
 flaxdiff/inputs/encoders.py,sha256=pjfbx4Rk7bLoE80MOfThZDm6YtsDncRekmn0Bmg_CwI,2963
 flaxdiff/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -62,7 +62,7 @@ flaxdiff/trainer/autoencoder_trainer.py,sha256=2FP2P-k9c0n_k3eT0trkq73dQrHRdBj9O
 flaxdiff/trainer/diffusion_trainer.py,sha256=reQEVWKTqKAeyCMQ-curPOfSRmBKxKooK8EVtUuorcM,14599
 flaxdiff/trainer/general_diffusion_trainer.py,sha256=BeDpJzgR8bUClJI4epQXlAul27MwiSfRW0lIBZSiPWk,28342
 flaxdiff/trainer/simple_trainer.py,sha256=raLS1shwpjJBT_bYXLAB2E4kA9MbwasDTzDTUqfCCUc,24312
-flaxdiff-0.2.6.1.dist-info/METADATA,sha256=zNcAo6k99Wb9N5gG8h685M1_1vC5Q_XEJ9lZpfkuBvM,24059
-flaxdiff-0.2.6.1.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
-flaxdiff-0.2.6.1.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.2.6.1.dist-info/RECORD,,
+flaxdiff-0.2.7.dist-info/METADATA,sha256=nwglJYeF2lH_MNq5PeFLR8TSPU-I9tzJUcBbTaLYxRM,24057
+flaxdiff-0.2.7.dist-info/WHEEL,sha256=wXxTzcEDnjrTwFYjLPcsW_7_XihufBwmpiBeiXNBGEA,91
+flaxdiff-0.2.7.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.2.7.dist-info/RECORD,,