demo

app.py

@@ -44,7 +44,7 @@ def main():
         # parse json to span attributes
         base_text_prompt, style_text_prompts, footnote_text_prompts, footnote_target_tokens,\
             color_text_prompts, color_names, color_rgbs, size_text_prompts_and_sizes, use_grad_guidance = parse_json(
-                json.loads(text_input))
+                json.loads(text_input), device)
 
         # create control input for region diffusion
         region_text_prompts, region_target_token_ids, base_tokens = get_region_diffusion_input(

utils/attention_utils.py

@@ -184,5 +184,5 @@ def get_token_maps(attention_maps, save_dir, width, height, obj_tokens, seed=0,
     token_maps_vis = plot_attention_maps([attention_maps_averaged, attention_maps_averaged_normalized],
                         obj_tokens, save_dir, seed, tokens_vis)
     attention_maps_averaged_normalized = [attn_mask.unsqueeze(1).repeat(
-        [1, 4, 1, 1]).cuda() for attn_mask in attention_maps_averaged_normalized]
+        [1, 4, 1, 1]).to(attention_maps_averaged_sum.device) for attn_mask in attention_maps_averaged_normalized]
     return attention_maps_averaged_normalized, token_maps_vis

utils/richtext_utils.py

@@ -27,7 +27,7 @@ def seed_everything(seed):
     torch.cuda.manual_seed(seed)
 
 
-def hex_to_rgb(hex_string, return_nearest_color=False):
+def hex_to_rgb(hex_string, return_nearest_color=False, device='cuda'):
     r"""
     Covert Hex triplet to RGB triplet.
     """
@@ -40,8 +40,8 @@ def hex_to_rgb(hex_string, return_nearest_color=False):
     rgb = torch.FloatTensor((red, green, blue))[None, :, None, None]/255.
     if return_nearest_color:
         nearest_color = find_nearest_color(rgb)
-        return rgb.cuda(), nearest_color
-    return rgb.cuda()
+        return rgb.to(device), nearest_color
+    return rgb.to(device)
 
 
 def find_nearest_color(rgb):
@@ -56,7 +56,7 @@ def find_nearest_color(rgb):
     return nearest_color
 
 
-def font2style(font):
+def font2style(font, device='cuda'):
     r"""
     Convert the font name to the style name.
     """
@@ -71,7 +71,7 @@ def font2style(font):
             'Akronim': 'Abstract Cubism, Pablo Picasso', }[font]
 
 
-def parse_json(json_str):
+def parse_json(json_str, device):
     r"""
     Convert the JSON string to attributes.
     """
@@ -121,7 +121,7 @@ def parse_json(json_str):
             if 'color' in span['attributes']:
                 use_grad_guidance = True
                 color_rgb, nearest_color = hex_to_rgb(
-                    span['attributes']['color'], True)
+                    span['attributes']['color'], True, device=device)
                 if prev_color_rgb == color_rgb:
                     prev_text_prompt = color_text_prompts[-1]
                     color_text_prompts[-1] = prev_text_prompt + \
@@ -197,8 +197,8 @@ def get_attention_control_input(model, base_tokens, size_text_prompts_and_sizes)
             word_pos.append(base_tokens.index(size_token)+1)
             font_sizes.append(font_size)
     if len(word_pos) > 0:
-        word_pos = torch.LongTensor(word_pos).cuda()
-        font_sizes = torch.FloatTensor(font_sizes).cuda()
+        word_pos = torch.LongTensor(word_pos).to(model.device)
+        font_sizes = torch.FloatTensor(font_sizes).to(model.device)
     else:
         word_pos = None
         font_sizes = None