Update app.py

app.py

@@ -34,6 +34,12 @@ from diffusers import (
 
 
 device = gr.State("cuda")
+unet = gr.State()
+vae = gr.State()
+text_encoder = gr.State()
+tokenizer = gr.State()
+noise_scheduler = gr.State()
+network = gr.State()
 
 pretrained_model_name_or_path = "stablediffusionapi/realistic-vision-v51" 
 revision = None
@@ -43,39 +49,35 @@ weight_dtype = torch.bfloat16
 pipe = StableDiffusionPipeline.from_pretrained("stablediffusionapi/realistic-vision-v51", 
                                             torch_dtype=torch.float16,safety_checker = None,
                                             requires_safety_checker = False).to(device.value)
-noise_scheduler = pipe.scheduler
+noise_scheduler.value = pipe.scheduler
 del pipe
-tokenizer = AutoTokenizer.from_pretrained(
+tokenizer.value = AutoTokenizer.from_pretrained(
         pretrained_model_name_or_path, subfolder="tokenizer", revision=revision
     )
-text_encoder = CLIPTextModel.from_pretrained(
+text_encoder.value = CLIPTextModel.from_pretrained(
         pretrained_model_name_or_path, subfolder="text_encoder", revision=revision
     )
-vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
-unet = UNet2DConditionModel.from_pretrained(
+vae.value = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae", revision=revision)
+unet.value = UNet2DConditionModel.from_pretrained(
         pretrained_model_name_or_path, subfolder="unet", revision=revision
     )
 
-unet.requires_grad_(False)
-unet.to(device.value, dtype=weight_dtype)
-vae.requires_grad_(False)
+unet.value.requires_grad_(False)
+unet.value.to(device.value, dtype=weight_dtype)
+vae.value.requires_grad_(False)
 
-text_encoder.requires_grad_(False)
-vae.requires_grad_(False)
-vae.to(device.value, dtype=weight_dtype)
-text_encoder.to(device.value, dtype=weight_dtype)
+text_encoder.value.requires_grad_(False)
+vae.value.requires_grad_(False)
+vae.value.to(device.value, dtype=weight_dtype)
+text_encoder.value.to(device.value, dtype=weight_dtype)
 print("")
 
 
 
-unet = gr.State(unet)
-vae = gr.State(vae)
-text_encoder = gr.State(text_encoder)
-tokenizer = gr.State(tokenizer)
-noise_scheduler = gr.State(noise_scheduler)
 
 
-network = gr.State()
+
+