Update shape.py

shape.py

@@ -8,4 +8,22 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 xm = load_model('transmitter', device=device)
 model = load_model('text300M', device=device)
 # model.load_state_dict(torch.load('./load/shapE_finetuned_with_330kdata.pth', map_location=device)['model_state_dict'])
-diffusion = diffusion_from_config_shape(load_config('diffusion'))
+diffusion = diffusion_from_config_shape(load_config('diffusion'))
+del xm
+del model
+del diffusion
+
+from diffusers import DDIMScheduler, DDPMScheduler, StableDiffusionPipeline
+
+
+pipe_kwargs = {
+    "tokenizer": None,
+    "safety_checker": None,
+    "feature_extractor": None,
+    "requires_safety_checker": False,
+    "torch_dtype": torch.float16 ,
+}
+self.pipe = StableDiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1-base",
+    **pipe_kwargs,
+)