Update app.py

app.py

@@ -69,7 +69,6 @@ pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-me
 
 pipe.load_lora_weights('ford442/sdxl-vae-bf16', weight_name='LoRA/bm-goth_epoch_9.safetensors')
 
-pipe.to(device=device, dtype=torch.bfloat16)
 
 #pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(torch.device("cuda:0"))
 #pipe = StableDiffusion3Pipeline.from_pretrained("ford442/RealVis_Medium_1.0b_bf16", torch_dtype=torch.bfloat16)
@@ -172,6 +171,15 @@ 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,