Update model.py

model.py

@@ -140,7 +140,7 @@ class StreamMultiDiffusion(nn.Module):
         self.i2t_processor = Blip2Processor.from_pretrained('Salesforce/blip2-opt-2.7b')
         self.i2t_model = Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-opt-2.7b')
 
-        self.pipe = load_model(model_key, self.sd_version, self.device, self.dtype)
+        self.pipe = load_model(model_key, self.sd_version, self.device, self.dtype).to(dtype=self.dtype)
 
         self.pipe.load_lora_weights(lora_key, weight_name=lora_weight_name, adapter_name='lcm')
         self.pipe.fuse_lora(
@@ -380,7 +380,7 @@ class StreamMultiDiffusion(nn.Module):
         """
         question = 'Question: What are in the image? Answer:'
         inputs = self.i2t_processor(image, question, return_tensors='pt')
-        out = self.i2t_model.generate(**{k: v.to(self.i2t_model.device) for k, v in inputs.items()}, max_new_tokens=77)
+        out = self.i2t_model.generate(**{k: v.to(self.i2t_model.device) for k, v in inputs.items()}) #, max_new_tokens=75)
         prompt = self.i2t_processor.decode(out[0], skip_special_tokens=True).strip()
         return prompt
 
@@ -1121,12 +1121,15 @@ class StreamMultiDiffusion(nn.Module):
         else:
             x_t_latent_plus_uc = x_t_latent  # (T * p, 4, h, w)
 
-        model_pred = self.pipe.unet(
-            x_t_latent_plus_uc.to(self.pipe.unet.dtype),  # (B, 4, h, w)
-            t_list,  # (B,)
-            encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
-            return_dict=False,
-        )[0]  # (B, 4, h, w)
+        try:
+            model_pred = self.unet(
+                x_t_latent_plus_uc.to(self.unet.dtype),  # (B, 4, h, w)
+                t_list,  # (B,)
+                encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
+                return_dict=False,
+            )[0]  # (B, 4, h, w)
+        except e:
+            print(e)
         print('222222222222222', model_pred.dtype)
 
         if self.bootstrap_steps[0] > 0: