update v1.1

app.py

@@ -33,6 +33,7 @@ from tools import BEN2
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--port', type=int, default=8080)
+parser.add_argument('--version', type=str, default='v1.1')
 parser.add_argument('--no_turbo', action='store_true')
 args = parser.parse_args()
 
@@ -169,17 +170,17 @@ def generate_image(
 
 _HEADER_ = '''
 <div style="text-align: center; max-width: 650px; margin: 0 auto;">
-    <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">DreamO</h1>
+    <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 1rem; display: contents;">DreamO v1.1</h1>
     <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.06.24: Updated to version 1.1 with significant improvements in image quality, reduced likelihood of body composition errors, and enhanced aesthetics. <a href='https://github.com/bytedance/DreamO/blob/main/dreamo_v1.1.md' target='_blank'>Learn more about this model</a>
 - 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
@@ -210,10 +211,10 @@ def create_demo():
                 width = gr.Slider(768, 1024, 1024, step=16, label="Width")
                 height = gr.Slider(768, 1024, 1024, step=16, label="Height")
                 num_steps = gr.Slider(8, 30, 12, step=1, label="Number of steps")
-                guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Guidance")
+                guidance = gr.Slider(1.0, 10.0, 4.5 if args.version == 'v1.1' else 3.5, step=0.1, label="Guidance")
                 seed = gr.Textbox(label="Seed (-1 for random)", value="-1")
+                ref_res = gr.Slider(512, 1024, 512, step=16, label="resolution for ref image, increase it if necessary")
                 with gr.Accordion("Advanced Options", open=False, visible=False):
-                    ref_res = gr.Slider(512, 1024, 512, step=16, label="resolution for ref image")
                     neg_prompt = gr.Textbox(label="Neg Prompt", value="")
                     neg_guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="Neg Guidance")
                     true_cfg = gr.Slider(1, 5, 1, step=0.1, label="true cfg")
@@ -304,7 +305,7 @@ def create_demo():
                     'id',
                     'ip',
                     'the woman wearing a dress, In the banquet hall',
-                    7698454872441022867,
+                    42,
                 ],
                 [
                     'example_inputs/dog1.png',
@@ -328,7 +329,7 @@ def create_demo():
                     'ip',
                     'ip',
                     'a man is dancing with a woman in the room',
-                    8303780338601106219,
+                    42,
                 ],
             ]
             gr.Examples(

dreamo/dreamo_pipeline.py

@@ -43,16 +43,26 @@ class DreamOPipeline(FluxPipeline):
         self.task_embedding = nn.Embedding(2, 3072)
         self.idx_embedding = nn.Embedding(10, 3072)
 
-    def load_dreamo_model(self, device, use_turbo=True):
+    def load_dreamo_model(self, device, use_turbo=True, version='v1.1'):
         # 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')
+        if version == 'v1':
+            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')
+            quality_lora_pos = load_file('models/dreamo_quality_lora_pos.safetensors')
+            quality_lora_neg = load_file('models/dreamo_quality_lora_neg.safetensors')
+        elif version == 'v1.1':
+            hf_hub_download(repo_id='ByteDance/DreamO', filename='v1.1/dreamo_sft_lora.safetensors', local_dir='models')
+            hf_hub_download(repo_id='ByteDance/DreamO', filename='v1.1/dreamo_dpo_lora.safetensors', local_dir='models')
+            sft_lora = load_file('models/v1.1/dreamo_sft_lora.safetensors')
+            dpo_lora = load_file('models/v1.1/dreamo_dpo_lora.safetensors')
+        else:
+            raise ValueError(f'there is no {version}')
         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:]
@@ -83,18 +93,27 @@ class DreamOPipeline(FluxPipeline):
             adapter_names.append('turbo')
             adapter_weights.append(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)
+        if version == 'v1':
+            # 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)
+        elif version == 'v1.1':
+            self.load_lora_weights(sft_lora, adapter_name='sft_lora')
+            adapter_names.append('sft_lora')
+            adapter_weights.append(1)
+            self.load_lora_weights(dpo_lora, adapter_name='dpo_lora')
+            adapter_names.append('dpo_lora')
+            adapter_weights.append(1.25)
 
         self.set_adapters(adapter_names, adapter_weights)
         self.fuse_lora(adapter_names=adapter_names, lora_scale=1)
+        self.unload_lora_weights()
 
         self.t5_embedding = self.t5_embedding.to(device)
         self.task_embedding = self.task_embedding.to(device)