Update app.py

app.py

@@ -147,7 +147,7 @@ class ModelWrapper:
 
         return output_image_list, f"Run successfully in {(end_time-start_time):.2f} seconds"
 
-
+@spaces.GPU()
 def get_x0_from_noise(sample, model_output, alphas_cumprod, timestep):
     alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)
     beta_prod_t = 1 - alpha_prod_t
@@ -159,18 +159,18 @@ class SDXLTextEncoder(torch.nn.Module):
     def __init__(self, model_id, revision, accelerator, dtype=torch.float32):
         super().__init__()
 
-        self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(accelerator.device).to(dtype=dtype)
-        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, subfolder="text_encoder_2", revision=revision).to(accelerator.device).to(dtype=dtype)
+        self.text_encoder_one = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", revision=revision).to(0).to(dtype=dtype)
+        self.text_encoder_two = CLIPTextModelWithProjection.from_pretrained(model_id, subfolder="text_encoder_2", revision=revision).to(0).to(dtype=dtype)
 
         self.accelerator = accelerator
 
     def forward(self, batch):
-        text_input_ids_one = batch['text_input_ids_one'].to(self.accelerator.device).squeeze(1)
-        text_input_ids_two = batch['text_input_ids_two'].to(self.accelerator.device).squeeze(1)
+        text_input_ids_one = batch['text_input_ids_one'].to(0).squeeze(1)
+        text_input_ids_two = batch['text_input_ids_two'].to(0).squeeze(1)
         prompt_embeds_list = []
 
         for text_input_ids, text_encoder in zip([text_input_ids_one, text_input_ids_two], [self.text_encoder_one, self.text_encoder_two]):
-            prompt_embeds = text_encoder(text_input_ids.to(text_encoder.device), output_hidden_states=True)
+            prompt_embeds = text_encoder(text_input_ids.to(0), output_hidden_states=True)
 
             pooled_prompt_embeds = prompt_embeds[0]
 
@@ -184,7 +184,7 @@ class SDXLTextEncoder(torch.nn.Module):
         
         return prompt_embeds, pooled_prompt_embeds
 
-@spaces.GPU(duration=100)
+
 def create_demo():
     TITLE = "# DMD2-SDXL Demo"
     model_id = "stabilityai/stable-diffusion-xl-base-1.0"