Update model.py

model.py

@@ -379,11 +379,8 @@ class StreamMultiDiffusion(nn.Module):
             A single string of text prompt.
         """
         question = 'Question: What are in the image? Answer:'
-        print(self.i2t_model.device)
         inputs = self.i2t_processor(image, question, return_tensors='pt')
-        print(inputs)
-        out = self.i2t_model.generate(**inputs, max_new_tokens=77)
-        print(out[0].device)
+        out = self.i2t_model.generate(**{k: v.to(self.i2t_model.device) for k, v in inputs.items()}, max_new_tokens=77)
         prompt = self.i2t_processor.decode(out[0], skip_special_tokens=True).strip()
         return prompt