a better model

app.py

@@ -28,19 +28,20 @@ from PIL import Image
 from torchvision.transforms.functional import normalize
 
 from dreamo.dreamo_pipeline import DreamOPipeline
-from dreamo.utils import img2tensor, resize_numpy_image_area, tensor2img
+from dreamo.utils import img2tensor, resize_numpy_image_area, tensor2img, resize_numpy_image_long
 from tools import BEN2
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--port', type=int, default=8080)
+parser.add_argument('--no_turbo', action='store_true')
 args = parser.parse_args()
 
 huggingface_hub.login(os.getenv('HF_TOKEN'))
 
-# try:
-#     shutil.rmtree('gradio_cached_examples')
-# except FileNotFoundError:
-#     print("cache folder not exist")
+try:
+    shutil.rmtree('gradio_cached_examples')
+except FileNotFoundError:
+    print("cache folder not exist")
 
 class Generator:
     def __init__(self):
@@ -63,7 +64,7 @@ class Generator:
         # load dreamo
         model_root = 'black-forest-labs/FLUX.1-dev'
         dreamo_pipeline = DreamOPipeline.from_pretrained(model_root, torch_dtype=torch.bfloat16)
-        dreamo_pipeline.load_dreamo_model(device, use_turbo=True)
+        dreamo_pipeline.load_dreamo_model(device, use_turbo=not args.no_turbo)
         self.dreamo_pipeline = dreamo_pipeline.to(device)
 
     @torch.no_grad()
@@ -126,10 +127,12 @@ def generate_image(
     for idx, (ref_image, ref_task) in enumerate(zip(ref_images, ref_tasks)):
         if ref_image is not None:
             if ref_task == "id":
+                ref_image = resize_numpy_image_long(ref_image, 1024)
                 ref_image = generator.get_align_face(ref_image)
             elif ref_task != "style":
                 ref_image = generator.bg_rm_model.inference(Image.fromarray(ref_image))
-            ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
+            if ref_task != "id":
+                ref_image = resize_numpy_image_area(np.array(ref_image), ref_res * ref_res)
             debug_images.append(ref_image)
             ref_image = img2tensor(ref_image, bgr2rgb=False).unsqueeze(0) / 255.0
             ref_image = 2 * ref_image - 1.0
@@ -170,9 +173,13 @@ _HEADER_ = '''
     <p style="font-size: 1rem; margin-bottom: 1.5rem;">Paper: <a href='https://arxiv.org/abs/2504.16915' target='_blank'>DreamO: A Unified Framework for Image Customization</a> | Codes: <a href='https://github.com/bytedance/DreamO' target='_blank'>GitHub</a></p>
 </div>
 
+🚩 Update Notes:
+- 2025.05.11: We have updated the model to mitigate over-saturation and plastic-face issues. The new version shows consistent improvements over the previous release.
+
 ❗️❗️❗️**User Guide:**
 - The most important thing to do first is to try the examples provided below the demo, which will help you better understand the capabilities of the DreamO model and the types of tasks it currently supports
 - For each input, please select the appropriate task type. For general objects, characters, or clothing, choose IP — we will remove the background from the input image. If you select ID, we will extract the face region from the input image (similar to PuLID). If you select Style, the background will be preserved, and you must prepend the prompt with the instruction: 'generate a same style image.' to activate the style task.
+- The most import hyperparameter in this demo is the guidance scale, which is set to 3.5 by default. If you notice that faces appear overly glossy or unrealistic—especially in ID tasks—you can lower the guidance scale (e.g., to 3). Conversely, if text rendering is poor or limb distortion occurs, increasing the guidance scale (e.g., to 4) may help.
 - To accelerate inference, we adopt FLUX-turbo LoRA, which reduces the sampling steps from 25 to 12 compared to FLUX-dev. Additionally, we distill a CFG LoRA, achieving nearly a twofold reduction in steps by eliminating the need for true CFG
 
 '''  # noqa E501
@@ -249,7 +256,7 @@ def create_demo():
                     'ip',
                     'ip',
                     'a purple toy holding a sign saying "DreamO", on the mountain',
-                    1563188099017016129,
+                    10441727852953907380,
                 ],
                 [
                     'example_inputs/perfume.png',

dreamo/dreamo_pipeline.py

@@ -44,24 +44,35 @@ class DreamOPipeline(FluxPipeline):
         self.idx_embedding = nn.Embedding(10, 3072)
 
     def load_dreamo_model(self, device, use_turbo=True):
+        # download models and load file
         hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo.safetensors', local_dir='models')
         hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo_cfg_distill.safetensors', local_dir='models')
+        hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo_quality_lora_pos.safetensors', local_dir='models')
+        hf_hub_download(repo_id='ByteDance/DreamO', filename='dreamo_quality_lora_neg.safetensors', local_dir='models')
         dreamo_lora = load_file('models/dreamo.safetensors')
         cfg_distill_lora = load_file('models/dreamo_cfg_distill.safetensors')
+        quality_lora_pos = load_file('models/dreamo_quality_lora_pos.safetensors')
+        quality_lora_neg = load_file('models/dreamo_quality_lora_neg.safetensors')
+
+        # load embedding
         self.t5_embedding.weight.data = dreamo_lora.pop('dreamo_t5_embedding.weight')[-10:]
         self.task_embedding.weight.data = dreamo_lora.pop('dreamo_task_embedding.weight')
         self.idx_embedding.weight.data = dreamo_lora.pop('dreamo_idx_embedding.weight')
         self._prepare_t5()
 
+        # main lora
         dreamo_diffuser_lora = convert_flux_lora_to_diffusers(dreamo_lora)
-        cfg_diffuser_lora = convert_flux_lora_to_diffusers(cfg_distill_lora)
         adapter_names = ['dreamo']
         adapter_weights = [1]
         self.load_lora_weights(dreamo_diffuser_lora, adapter_name='dreamo')
-        if cfg_diffuser_lora is not None:
-            self.load_lora_weights(cfg_diffuser_lora, adapter_name='cfg')
-            adapter_names.append('cfg')
-            adapter_weights.append(1)
+
+        # cfg lora to avoid true image cfg
+        cfg_diffuser_lora = convert_flux_lora_to_diffusers(cfg_distill_lora)
+        self.load_lora_weights(cfg_diffuser_lora, adapter_name='cfg')
+        adapter_names.append('cfg')
+        adapter_weights.append(1)
+
+        # turbo lora to speed up (from 25+ step to 12 step)
         if use_turbo:
             self.load_lora_weights(
                 hf_hub_download(
@@ -72,7 +83,18 @@ class DreamOPipeline(FluxPipeline):
             adapter_names.append('turbo')
             adapter_weights.append(1)
 
-        self.fuse_lora(adapter_names=adapter_names, adapter_weights=adapter_weights, lora_scale=1)
+        # quality loras, one pos, one neg
+        quality_lora_pos = convert_flux_lora_to_diffusers(quality_lora_pos)
+        self.load_lora_weights(quality_lora_pos, adapter_name='quality_pos')
+        adapter_names.append('quality_pos')
+        adapter_weights.append(0.15)
+        quality_lora_neg = convert_flux_lora_to_diffusers(quality_lora_neg)
+        self.load_lora_weights(quality_lora_neg, adapter_name='quality_neg')
+        adapter_names.append('quality_neg')
+        adapter_weights.append(-0.8)
+
+        self.set_adapters(adapter_names, adapter_weights)
+        self.fuse_lora(adapter_names=adapter_names, lora_scale=1)
 
         self.t5_embedding = self.t5_embedding.to(device)
         self.task_embedding = self.task_embedding.to(device)

dreamo/utils.py

@@ -117,6 +117,16 @@ def resize_numpy_image_area(image, area=512 * 512):
     image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
     return image
 
+def resize_numpy_image_long(image, long_edge=768):
+    h, w = image.shape[:2]
+    if max(h, w) <= long_edge:
+        return image
+    k = long_edge / max(h, w)
+    h = int(h * k)
+    w = int(w * k)
+    image = cv2.resize(image, (w, h), interpolation=cv2.INTER_AREA)
+    return image
+
 
 # reference: https://github.com/huggingface/diffusers/pull/9295/files
 def convert_flux_lora_to_diffusers(old_state_dict):