Update app.py

app.py

@@ -20,20 +20,7 @@ from huggingface_hub import snapshot_download
 import spaces
 
 
-generator = gr.State()
-unet = gr.State()
-vae = gr.State()
-text_encoder = gr.State()
-tokenizer = gr.State()
-noise_scheduler = gr.State()
-network = gr.State()
-#device = gr.State(torch.device("cuda"))
-device = "cuda"
-#generator = torch.Generator(device=device)
-young = gr.State()
-pointy = gr.State()
-wavy = gr.State()
-thick = gr.State()
+
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
@@ -45,7 +32,7 @@ df = torch.load(f"{models_path}/files/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/files/weight_dimensions.pt")
 pinverse = torch.load(f"{models_path}/files/pinverse_1000pc.pt", map_location=torch.device('cpu')).bfloat16().to(device)
 
-unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
+
 
 
 
@@ -408,6 +395,29 @@ intro = """
 
 
 with gr.Blocks(css="style.css") as demo:
+    generator = gr.State()
+    unet = gr.State()
+    vae = gr.State()
+    text_encoder = gr.State()
+    tokenizer = gr.State()
+    noise_scheduler = gr.State()
+    network = gr.State()
+    #device = gr.State(torch.device("cuda"))
+    device = "cuda"
+    #generator = torch.Generator(device=device)
+    young = gr.State()
+    pointy = gr.State()
+    wavy = gr.State()
+    thick = gr.State()
+    unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
+
+
+
+
+
+
+
+    
     gr.HTML(intro)
     
     gr.Markdown("""<div style="text-align: justify;"> In this demo, you can get an identity-encoding model by sampling or inverting. To use a model previously downloaded from this demo see \"Uploading a model\" in the Advanced Options. Next, you can generate new images from it, or edit the identity encoded in the model and generate images from the edited model. We provide detailed instructions and tips at the bottom of the page.""")