Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import numpy as np
 import random
+import gc
 import json
 import torch
 import spaces
@@ -27,15 +28,6 @@ if torch.cuda.is_available():
 else:
     torch_dtype = torch.float32
 
-# Initialize models from base SD3.5
-vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
-text_encoder = CLIPTextModelWithProjection.from_pretrained(model_repo_id, subfolder="text_encoder")
-text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(model_repo_id, subfolder="text_encoder_2")
-text_encoder_3 = T5EncoderModel.from_pretrained(mdoel_repo_id, subfolder="text_encoder_3")
-tokenizer = CLIPTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
-tokenizer_2 = CLIPTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer_2")
-tokenizer_3 = T5Tokenizer.from_pretrained(model_repo_id, subfolder="tokenizer_3")
-
 # Initialize transformer
 config_file = hf_hub_download(repo_id=model_repo_id, filename="transformer/config.json")
 with open(config_file, "r") as fp:
@@ -53,6 +45,20 @@ with safe_open(model_file, framework="pt") as f:
 state_dict = convert_sd3_transformer_checkpoint_to_diffusers(state_dict)
 transformer.load_state_dict(state_dict)
 
+# Try to keep memory usage down
+del state_dict
+gc.collect()
+
+# Initialize models from base SD3.5
+vae = AutoencoderKL.from_pretrained(model_repo_id, subfolder="vae")
+text_encoder = CLIPTextModelWithProjection.from_pretrained(model_repo_id, subfolder="text_encoder")
+text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(model_repo_id, subfolder="text_encoder_2")
+text_encoder_3 = T5EncoderModel.from_pretrained(model_repo_id, subfolder="text_encoder_3")
+tokenizer = CLIPTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer")
+tokenizer_2 = CLIPTokenizer.from_pretrained(model_repo_id, subfolder="tokenizer_2")
+tokenizer_3 = T5Tokenizer.from_pretrained(model_repo_id, subfolder="tokenizer_3")
+
+
 # Create pipeline from our models
 pipe = StableDiffusion3Pipeline(
     vae=vae,