Merge pull request #105 from BitMind-AI/expand-models/image-finetuned

Adding two finetuned image models to expand validator challenges

bitmind/constants.py

@@ -89,6 +89,18 @@
             },
             "enable_cpu_offload": False,
             "pipeline": "FluxPipeline"
+        },
+        {
+            "path": "prompthero/openjourney-v4",
+            "use_safetensors": True,
+            "torch_dtype": torch.float16,
+            "pipeline": "StableDiffusionPipeline"
+        },
+        {
+            "path": "cagliostrolab/animagine-xl-3.1",
+            "use_safetensors": True,
+            "torch_dtype": torch.float16,
+            "pipeline": "StableDiffusionXLPipeline"
         }
     ]
 }

bitmind/synthetic_image_generation/synthetic_image_generator.py

@@ -1,6 +1,6 @@
 from transformers import pipeline
 from transformers import set_seed
-from diffusers import StableDiffusionXLPipeline, FluxPipeline
+from diffusers import StableDiffusionXLPipeline, FluxPipeline, StableDiffusionPipeline
 import bittensor as bt
 import numpy as np
 import torch
@@ -260,7 +260,7 @@ def get_tokenizer_with_min_len(self):
             tuple: A tuple containing the tokenizer and its maximum token length.
         """
         # Check if a second tokenizer is available in the diffuser
-        if self.diffuser.tokenizer_2:
+        if hasattr(self.diffuser, 'tokenizer_2'):
             if self.diffuser.tokenizer.model_max_length > self.diffuser.tokenizer_2.model_max_length:
                 return self.diffuser.tokenizer_2, self.diffuser.tokenizer_2.model_max_length
         return self.diffuser.tokenizer, self.diffuser.tokenizer.model_max_length