32

artist_lib.py

@@ -88,6 +88,7 @@ def generate_tone(note, octave, duration):
 
 def draw(inp, this_model, force_new):
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = "torch.float16" if torch.cuda.is_available() else "torch.float32"
     drawing = inp
     if this_model == "stable-diffusion-2":
         this_model_addr = "stabilityai/stable-diffusion-2"
@@ -110,7 +111,7 @@ def draw(inp, this_model, force_new):
             print("found drawing ", drawing_filename)
             return Image.open(drawing_filename) 
     print("generating drawing '", drawing, "'", drawing_filename)
-    pipe = StableDiffusionPipeline.from_pretrained(this_model_addr, torch_dtype=torch.float16)
+    pipe = StableDiffusionPipeline.from_pretrained(this_model_addr, torch_dtype=dtype)
     pipe.enable_attention_slicing()
     pipe = pipe.to(device)
     image = pipe(drawing).images[0]  
@@ -140,10 +141,11 @@ def write_blog(inp, this_model, min_length, max_length, force_new):
               return file.read()
     print("generating blog '", blog_post_name, "'", target_filename)
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    #generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device=device, torch_dtype=torch.float16)
-    #generator = pipeline('text-generation', model=this_model_addr, torch_dtype=torch.float16)
+    dtype = "torch.float16" if torch.cuda.is_available() else "torch.float32"
+    #generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B', device=device, torch_dtype=dtype)
+    #generator = pipeline('text-generation', model=this_model_addr, torch_dtype=dtype)
     #generator = pipeline('text-generation', model=this_model_addr)
-    generator = pipeline('text-generation', model=this_model_addr, device=device, torch_dtype=torch.float16)
+    generator = pipeline('text-generation', model=this_model_addr, device=device, torch_dtype=dtype)
     # AttributeError: 'TextGenerationPipeline' object has no attribute 'enable_attention_slicing'
     #generator.enable_attention_slicing()
     res = generator(blog_post_name, min_length=min_length, max_length=max_length, do_sample=True, temperature=0.7)