Update app.py

app.py

@@ -12,14 +12,15 @@ from utils import load_models, save_model_w2w, save_model_for_diffusers
 from sampling import sample_weights
 from huggingface_hub import snapshot_download
 
-#global device
-#global generator 
-#global unet
-#global vae 
-#global text_encoder
-#global tokenizer
-#global noise_scheduler
+global device
+global generator 
+global unet
+global vae 
+global text_encoder
+global tokenizer
+global noise_scheduler
 device = "cuda:0"
+generator = torch.Generator(device=device)
 
 models_path = snapshot_download(repo_id="Snapchat/w2w")
 
@@ -31,30 +32,33 @@ df = torch.load(f"{models_path}/identity_df.pt")
 weight_dimensions = torch.load(f"{models_path}/weight_dimensions.pt")
 
 unet, vae, text_encoder, tokenizer, noise_scheduler = load_models(device)
-network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
-#global network
-
-#def sample_model():
-#    global unet
-#    del unet
-#    global network
-#    unet, _, _, _, _ = load_models(device)
-
-def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
-    #global device
-    #global generator 
-    #global unet
-    #global vae 
-    #global text_encoder
-    #global tokenizer
-    #global noise_scheduler
-    generator = torch.Generator(device=device).manual_seed(seed)
+global network
+
+def sample_model():
+    global unet
+    del unet
+    global network
+    unet, _, _, _, _ = load_models(device)
+    network = sample_weights(unet, proj, mean, std, v[:, :1000], device, factor = 1.00)
+ 
+
+@torch.no_grad()
+def inference( prompt, negative_prompt, guidance_scale, ddim_steps, seed):
+    global device
+    global generator 
+    global unet
+    global vae 
+    global text_encoder
+    global tokenizer
+    global noise_scheduler
+    generator = generator.manual_seed(seed)
     latents = torch.randn(
         (1, unet.in_channels, 512 // 8, 512 // 8),
         generator = generator,
         device = device
     ).bfloat16()
    
+
     text_input = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
 
     text_embeddings = text_encoder(text_input.input_ids.to(device))[0]
@@ -87,8 +91,12 @@ def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
 
     return [image] 
 
-with gr.Blocks() as demo:
+    
+
+css = ''
+with gr.Blocks(css=css) as demo:
     gr.Markdown("# <em>weights2weights</em> Demo")
+    gr.Markdown("Demo for the [h94/IP-Adapter-FaceID model](https://huggingface.co/h94/IP-Adapter-FaceID) - Generate AI images with your own face - Non-commercial license")
     with gr.Row():
         with gr.Column():
             files = gr.Files(
@@ -106,9 +114,9 @@ with gr.Blocks() as demo:
                        placeholder="sks person", 
                        value="sks person")
             negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, cartoon", value="low quality, blurry, unfinished, cartoon")
-            seed = gr.Number(value=5, label="Seed", interactive=True)
+            seed = gr.Number(value=5, precision=0, label="Seed", interactive=True)
             cfg = gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
-            steps = gr.Slider(label="Inference Steps", value=50, step=1, minimum=0, maximum=100, interactive=True)
+            steps = gr.Slider(label="Inference Steps", precision=0, value=50, step=1, minimum=0, maximum=100, interactive=True)
 
 
             submit = gr.Button("Submit")
@@ -116,10 +124,15 @@ with gr.Blocks() as demo:
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
 
-        #sample.click(fn=sample_model)
+        sample.click(fn=sample_model)
         
         submit.click(fn=inference,
                     inputs=[prompt, negative_prompt, cfg, steps, seed],
                     outputs=gallery)
             
-demo.launch(share=True)
+
+
+
+            
+            
+demo.launch(share=True)