Update app.py

app.py

@@ -147,7 +147,16 @@ def infer(
     #upscaler_2.to(torch.device('cpu'))
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device='cpu').manual_seed(seed)
+    generator = torch.Generator(device='cuda').manual_seed(seed)
+    if expanded_only:
+        pipe.to('cpu')
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+    else:
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        pipe.to(device=device, dtype=torch.bfloat16)
+        gc.collect()
     if expanded:
         system_prompt_rewrite = (
         "You are an AI assistant that rewrites image prompts to be more descriptive and detailed."
@@ -171,15 +180,7 @@ def infer(
         attention_mask_2 = encoded_inputs_2["attention_mask"].to(device)
         print("-- tokenize prompt --")
           # Google T5
-        if expanded_only:
-            pipe.to('cpu')
-            torch.cuda.empty_cache()
-            torch.cuda.reset_peak_memory_stats()
-        else:
-            torch.cuda.empty_cache()
-            torch.cuda.reset_peak_memory_stats()
-            pipe.to(device=device, dtype=torch.bfloat16)
-            gc.collect()
+
         #input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
         outputs = model.generate(
             input_ids=input_ids,