Fix: SES

Author: yjy415

diffsynth/utils/inference_time_scaling/ses.py

@@ -68,10 +68,9 @@ def _load_model(self):
 
     def get_score(self, image_pil, text_prompt):
         try:
-            with torch.no_grad():
+            with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float32):
                 if self.reward_name == "pick":
                     inputs = self.processor(text=[text_prompt], images=[image_pil], return_tensors="pt", padding="max_length", truncation=True, max_length=77).to(self.device)
-                    inputs['pixel_values'] = inputs['pixel_values'].to(self.dtype)
                     outputs = self.model(**inputs)
                     return outputs.logits_per_image[0, 0].item()
 
@@ -99,14 +98,13 @@ def run_ses_cem(
     popsize=10,
     k_elites=5,
     wavelet_name="db1",
-    dwt_level=5,
-    lambda_prior=1e-3
+    dwt_level=4,
 ):
     latent_h, latent_w = base_latents.shape[-2], base_latents.shape[-1]
     c_low_init, c_high_fixed_batch = split_dwt(base_latents, wavelet_name, dwt_level)
     c_high_fixed = c_high_fixed_batch[0]    
-    c_low_shape = c_low_init.shape[1:]     
-    mu = c_low_init.view(-1).cpu() 
+    c_low_shape = c_low_init.shape[1:]
+    mu = torch.zeros_like(c_low_init.view(-1).cpu()) 
     sigma_sq = torch.ones_like(mu) * 1.0 
 
     best_overall = {"fitness": -float('inf'), "score": -float('inf'), "c_low": c_low_init[0]}
@@ -135,16 +133,12 @@ def run_ses_cem(
             img = pipeline_callback(z_recon)
 
             score = scorer.get_score(img, prompt)
-            penalty = lambda_prior * (torch.norm(c_low_sample.float())**2).item()
-            fitness = score - penalty
-            
             res = {
-                "fitness": fitness, 
                 "score": score, 
                 "c_low": c_low_sample.cpu()
             }
             batch_results.append(res)
-            if fitness > best_overall['fitness']:
+            if score > best_overall['score']:
                 best_overall = res
 
             eval_count += 1
@@ -156,7 +150,7 @@ def run_ses_cem(
 
         if not batch_results: break
         elite_db.extend(batch_results)        
-        elite_db.sort(key=lambda x: x['fitness'], reverse=True)        
+        elite_db.sort(key=lambda x: x['score'], reverse=True)        
         elite_db = elite_db[:k_elites]        
         elites_flat = torch.stack([x['c_low'].view(-1) for x in elite_db])
         mu_new = torch.mean(elites_flat, dim=0)

examples/flux/model_inference/FLUX.1-dev-SES.py

@@ -12,7 +12,7 @@
     ],
 )
 
-prompt = "A solo girl with silver wavy hair and blue eyes, wearing a blue dress, underwater, air bubbles, floating hair."
+prompt = "A magical forest where trees are made of candy"
 negative_prompt = "nsfw, low quality"
 
 image = pipe(

examples/flux2/model_inference/FLUX.2-dev-SES.py

@@ -21,7 +21,7 @@
     ],
     tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"),
 )
-prompt = "A hermit crab using a soda can as its shell on the beach. The can has the text 'BFL Diffusers' on it."
+prompt = "A magical forest where trees are made of candy"
 
 image = pipe(
     prompt, 
@@ -31,7 +31,7 @@
     enable_ses=True,
     ses_reward_model="pick",
     ses_eval_budget=20,
-    ses_inference_steps=10
+    ses_inference_steps=20
 )
 
 image.save("image_FLUX.2-dev_ses.jpg")

examples/qwen_image/model_inference/Qwen-Image-SES.py

@@ -12,15 +12,15 @@
     tokenizer_config=ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
 )
 
-prompt = "水下少女，身穿蓝裙，周围有气泡。"
+prompt = "一把精致的汉服折扇，上面绘有山水"
 
 image = pipe(
     prompt, 
     seed=0, 
     num_inference_steps=40,
     enable_ses=True,
     ses_reward_model="pick",
-    ses_eval_budget=20,
+    ses_eval_budget=30,
     ses_inference_steps=10
 )
 

examples/z_image/model_inference/Z-Image-SES.py

@@ -12,7 +12,7 @@
     ],
     tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
 )
-prompt = "Chinese woman in red Hanfu holding a fan, with a bright yellow neon lightning bolt floating above her palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
+prompt = "A cyberpunk girl with neon glowing eyes"
 
 image = pipe(
     prompt=prompt, 
@@ -25,6 +25,4 @@
     ses_eval_budget=20,
     ses_inference_steps=10
 )
-image.save("image_Z-Image_ses.jpg")
-
-
+image.save("image_Z-Image_ses.jpg")

examples/z_image/model_inference/Z-Image-Turbo-SES.py

@@ -12,15 +12,15 @@
     tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
 )
 
-prompt = "Chinese woman in red Hanfu holding a fan, with a bright yellow neon lightning bolt floating above her palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
+prompt = "A cyberpunk girl with neon glowing eyes"
 
 image = pipe(
     prompt=prompt, 
     seed=42, 
     rand_device="cuda",
     enable_ses=True,
     ses_reward_model="pick",
-    ses_eval_budget=50,
+    ses_eval_budget=30,
     ses_inference_steps=8
 )
 image.save("image_Z-Image-Turbo_ses.jpg")