add dalle-mini

hydra_node/config.py

@@ -10,7 +10,9 @@ from dotmap import DotMap
 from icecream import ic
 from sentry_sdk import capture_exception
 from sentry_sdk.integrations.threading import ThreadingIntegration
-from hydra_node.models import StableDiffusionModel
+from hydra_node.models import StableDiffusionModel, DalleMiniModel
+
+model_map = {"stable_diffusion": StableDiffusionModel, "dalle-mini": DalleMiniModel}
 
 def no_init(loading_code):
     def dummy(self):
@@ -97,7 +99,7 @@ def init_config_model():
     load_time = time.time()
 
     try:
-        model = no_init(lambda: StableDiffusionModel(config))
+        model = no_init(lambda: model_map[config.model_name](config))
     except Exception as e:
         logger.error(f"Failed to load model: {str(e)}")
         capture_exception(e)

hydra_node/models.py

@@ -110,4 +110,44 @@ class StableDiffusionModel(nn.Module):
             x_sample = np.ascontiguousarray(x_sample)
             images.append(x_sample) 
 
-        return images
\ No newline at end of file
+        return images
+
+    @torch.no_grad()
+    def sample_from_image(self, request):
+        return
+
+class DalleMiniModel(nn.Module):
+    def __init__(self, config):
+        nn.Module.__init__(self)
+        from min_dalle import MinDalle
+
+        self.config = config
+        self.model = MinDalle(
+            models_root=config.model_path,
+            dtype=torch.float16,
+            device='cuda',
+            is_mega=True, 
+            is_reusable=True
+        )
+
+    @torch.no_grad()
+    def sample(self, request):
+        if request.seed is not None:
+            seed = request.seed
+        else:
+            seed = -1
+
+        images = self.model.generate_imagS(
+            text=request.prompt,
+            seed=seed,
+            grid_size=4,
+            is_seamless=False,
+            temperature=request.temp,
+            top_k=request.top_k,
+            supercondition_factor=request.scale,
+            is_verbose=False
+        )
+        images = images.to('cpu').numpy().transpose(0, 2, 3, 1)
+        images = np.ascontiguousarray(images)
+        return images
+

main.py

@@ -63,6 +63,8 @@ class GenerationRequest(BaseModel):
     scale: float = 7.0
     dynamic_threshold: float = None
     seed: int = None
+    temp: float = 1.0
+    top_k: int = 256
 
 class GenerationOutput(BaseModel):
     output: List[str]