update

app.py

@@ -81,7 +81,7 @@ def initialize_models():
     return generator, id_model, pose_model, quality_model
 
 @spaces.GPU
-def image_generation(input_image, quality, use_target_pose, pose, dimension, progress=gr.Progress()):
+def image_generation(input_image, quality, random_perturbation, sigma, dimension, progress=gr.Progress()):
     generator, id_model, pose_model, quality_model = initialize_models()
 
     generated_images = []
@@ -93,7 +93,7 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension, pro
         input_image.div_(255).sub_(0.5).div_(0.5)
         feature = id_model(input_image).clone().detach().cpu().numpy()
 
-    if not use_target_pose:
+    if not random_perturbation:
         features = []
         norm = np.linalg.norm(feature, 2, 1, True)
         for i in progress.tqdm(np.arange(0, 4.8, 2), desc="Generating images"):
@@ -102,16 +102,15 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension, pro
             updated_feature = updated_feature / np.linalg.norm(updated_feature, 2, 1, True) * norm
             features.append(updated_feature)
         features = torch.tensor(np.vstack(features)).float().to(device)
-        if quality > 25:
+        if quality > 22:
             images, _ = generator.gen_image(features, quality_model, id_model, q_target=quality)
         else:
             _, _, images, *_ = generator(features)
     else:
         features = torch.repeat_interleave(torch.tensor(feature), 3, dim=0)
-        features = sample_nearby_vectors(features, [0.7], [1]).float().to(device)
-        if quality > 25 or pose > 20:
-            images, _ = generator.gen_image(features, quality_model, id_model, pose_model=pose_model,
-                                            q_target=quality, pose=pose, class_rep=features)
+        features = sample_nearby_vectors(features, [sigma], [1]).float().to(device)
+        if quality > 22:
+            images, _ = generator.gen_image(features, quality_model, id_model, q_target=quality, class_rep=features)
         else:
             _, _, images, *_ = generator(features)
 
@@ -122,7 +121,7 @@ def image_generation(input_image, quality, use_target_pose, pose, dimension, pro
     return generated_images
 
 @spaces.GPU
-def process_input(image_input, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose, progress=gr.Progress()):
+def process_input(image_input, num1, num2, num3, num4, num5, num6, num7, num8, random_seed, target_quality, random_perturbation, sigma, progress=gr.Progress()):
     # Ensure all dimension numbers are within [0, 512)
     num1, num2, num3, num4 = [max(0, min(int(n), 511)) for n in [num1, num2, num3, num4]]
 
@@ -136,7 +135,7 @@ def process_input(image_input, num1, num2, num3, num4, random_seed, target_quali
         input_data = Image.open(image_input)
         input_data = np.array(input_data.resize((112, 112)))
 
-    generated_images = image_generation(input_data, target_quality, use_target_pose, target_pose, [num1, num2, num3, num4], progress)
+    generated_images = image_generation(input_data, target_quality, random_perturbation, sigma, [num1, num2, num3, num4, num5, num6, num7, num8], progress)
 
     return generated_images
 
@@ -148,11 +147,15 @@ def select_image(value, images):
 
 def toggle_inputs(use_pose):
     return [
-        gr.update(visible=use_pose, interactive=use_pose),  # target_pose
+        gr.update(visible=use_pose, interactive=use_pose),  # sigma
         gr.update(interactive=not use_pose),  # num1
         gr.update(interactive=not use_pose),  # num2
         gr.update(interactive=not use_pose),  # num3
         gr.update(interactive=not use_pose),  # num4
+        gr.update(interactive=not use_pose),  # num5
+        gr.update(interactive=not use_pose),  # num6
+        gr.update(interactive=not use_pose),  # num7
+        gr.update(interactive=not use_pose),  # num8
     ]
 
 
@@ -186,16 +189,20 @@ def main():
 
                 with gr.Row():
                     num1 = gr.Number(label="Dimension 1", value=0, minimum=0, maximum=511, step=1)
-                    num2 = gr.Number(label="Dimension 2", value=50, minimum=0, maximum=511, step=1)
-                    num3 = gr.Number(label="Dimension 3", value=100, minimum=0, maximum=511, step=1)
-                    num4 = gr.Number(label="Dimension 4", value=200, minimum=0, maximum=511, step=1)
+                    num2 = gr.Number(label="Dimension 2", value=25, minimum=0, maximum=511, step=1)
+                    num3 = gr.Number(label="Dimension 3", value=50, minimum=0, maximum=511, step=1)
+                    num4 = gr.Number(label="Dimension 4", value=75, minimum=0, maximum=511, step=1)
+                    num5 = gr.Number(label="Dimension 5", value=100, minimum=0, maximum=511, step=1)
+                    num6 = gr.Number(label="Dimension 6", value=125, minimum=0, maximum=511, step=1)
+                    num7 = gr.Number(label="Dimension 7", value=150, minimum=0, maximum=511, step=1)
+                    num8 = gr.Number(label="Dimension 8", value=200, minimum=0, maximum=511, step=1)
 
                 random_seed = gr.Number(label="Random Seed", value=42, minimum=0, maximum=MAX_SEED, step=1)
                 target_quality = gr.Slider(label="Minimum Quality", minimum=22, maximum=35, step=1, value=24)
 
                 with gr.Row():
-                    use_target_pose = gr.Checkbox(label="Use Target Pose")
-                    target_pose = gr.Slider(label="Target Pose", value=0, minimum=0, maximum=90, step=1, visible=False)
+                    random_perturbation = gr.Checkbox(label="Random Perturbation")
+                    sigma = gr.Slider(label="Sigma value", value=0, minimum=0, maximum=1, step=0.1, visible=False)
 
                 submit = gr.Button("Submit", variant="primary")
 
@@ -219,10 +226,10 @@ def main():
                             - These values are added to the dimensions (before normalization), **please ignore it if pose editing is on**.
                             """)
 
-        use_target_pose.change(
+        random_perturbation.change(
             fn=toggle_inputs,
-            inputs=[use_target_pose],
-            outputs=[target_pose, num1, num2, num3, num4]
+            inputs=[random_perturbation],
+            outputs=[sigma, num1, num2, num3, num4, num5, num6, num7, num8]
         )
 
         generated_images = gr.State([])
@@ -237,7 +244,7 @@ def main():
             outputs=[generation_time]
         ).then(
             fn=process_input,
-            inputs=[image_file, num1, num2, num3, num4, random_seed, target_quality, use_target_pose, target_pose],
+            inputs=[image_file, num1, num2, num3, num4, num5, num6, num7, num8, random_seed, target_quality, random_perturbation, sigma],
             outputs=[generated_images]
         ).then(
             fn=done,