Adding support for FastComposer

add device for running with cuda by default
open-mmlab · Sep 8, 2023 · 4bbacbb · 4bbacbb
1 parent 312c5f0
commit 4bbacbb
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/configs/fastcomposer/README.md b/configs/fastcomposer/README.md
@@ -84,7 +84,10 @@ if num_subject_in_text != len(image):
 if seed == -1:
     seed = np.random.randint(0, 1000000)
 
-generator = torch.manual_seed(seed)
+device = torch.device('cuda' if torch.cuda.is_available(
+    ) else 'cpu')
+generator = torch.Generator(device=device)
+generator.manual_seed(seed)
 
 output_dict = fastcomposer.infer(prompt,
                                  negative_prompt=negative_prompt,

diff --git a/configs/fastcomposer/README_zh-CN.md b/configs/fastcomposer/README_zh-CN.md
@@ -85,7 +85,10 @@ if num_subject_in_text != len(image):
 if seed == -1:
     seed = np.random.randint(0, 1000000)
 
-generator = torch.manual_seed(seed)
+device = torch.device('cuda' if torch.cuda.is_available(
+    ) else 'cpu')
+generator = torch.Generator(device=device)
+generator.manual_seed(seed)
 
 output_dict = fastcomposer.infer(prompt,
                                  negative_prompt=negative_prompt,

diff --git a/demo/gradio_fastcomposer.py b/demo/gradio_fastcomposer.py
@@ -20,6 +20,9 @@ class ModelWrapper:
     def __init__(self, model):
         super().__init__()
         self.model = model
+        self.device = torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = self.model.to(self.device)
 
     def inference(
         self,
@@ -58,7 +61,8 @@ def inference(
         if seed == -1:
             seed = np.random.randint(0, 1000000)
 
-        generator = torch.manual_seed(seed)
+        generator = torch.Generator(device=self.device)
+        generator.manual_seed(seed)
 
         return (
             self.model.infer(