generation

app.py

@@ -11,7 +11,7 @@ from transformers import AutoTokenizer, AutoModel
 from models.gen_pipeline import NextStepPipeline
 from utils.aspect_ratio import center_crop_arr_with_buckets
 
-HF_HUB = "stepfun-ai/NextStep-1-Large-Edit"
+HF_HUB = "stepfun-ai/NextStep-1-Large"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained(HF_HUB, local_files_only=False, trust_remote_code=True)
@@ -30,15 +30,15 @@ def infer(
     ref=None,
     seed=0,
     text_cfg=7.5,
-    img_cfg=2.0,
+    img_cfg=1.0,
     num_inference_steps=30,
     positive_prompt=DEFAULT_POSITIVE_PROMPT,
     negative_prompt=DEFAULT_NEGATIVE_PROMPT,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if ref is None:
-        gr.Warning("⚠️ Please upload an image!")
-        return None
+    #if ref is None:
+    #    gr.Warning("⚠️ Please upload an image!")
+    #    return None
 
     if prompt in [None, ""]:
         gr.Warning("⚠️ Please enter a prompt!")
@@ -55,8 +55,7 @@ def infer(
 
     image = pipeline.generate_image(
         captions=editing_caption,
-        images=input_image,
-        num_images_per_caption=2,
+        num_images_per_caption=1,
         positive_prompt=positive_prompt,
         negative_prompt=negative_prompt,
         hw=(input_image.size[1], input_image.size[0]),
@@ -69,15 +68,7 @@ def infer(
         progress=True,
     )
 
-    return image[0], image[1]
-
-
-examples = [
-    ["Modify the image to make the white horse run towards the camera.", "assets/1.jpg"],
-    ["Change the background to the sea view.", "assets/2.jpg"],
-    ["Add a pirate hat to the dog's head. Change the background to a stormy sea with dark clouds. Include the text 'NextStep-Edit' in bold white letters at the top portion of the image.", "assets/3.jpg"],
-    ["Change to Ghibli style.", "assets/4.jpg"],
-]
+    return image[0]
 
 css = """
 #col-container {
@@ -102,13 +93,13 @@ with gr.Blocks(css=css) as demo:
             run_button = gr.Button("Run", scale=0, variant="primary")
 
         with gr.Row():
-            ref = gr.Image(label="Reference Image", show_label=True, type="pil", height=400)
+            #ref = gr.Image(label="Reference Image", show_label=True, type="pil", height=400)
 
             with gr.Accordion("Advanced Settings", open=True):
                 positive_prompt = gr.Text(
                     label="Positive Prompt",
                     show_label=False,
-                    max_lines=2,
+                    max_lines=1,
                     placeholder="Enter your positive prompt",
                     container=False,
                 )
@@ -153,9 +144,9 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Row():
             result_1 = gr.Image(label="Result 1", show_label=False, container=True, height=400, visible=False)
-            result_2 = gr.Image(label="Result 2", show_label=False, container=True, height=400, visible=False)
+            #result_2 = gr.Image(label="Result 2", show_label=False, container=True, height=400, visible=False)
 
-        gr.Examples(examples=examples, inputs=[prompt, ref])
+        #gr.Examples(examples=examples, inputs=[prompt, ref])
 
     def show_result():
         return gr.update(visible=True), gr.update(visible=True)
@@ -165,7 +156,7 @@ with gr.Blocks(css=css) as demo:
         fn=infer,
         inputs=[
             prompt,
-            ref,
+            #ref,
             seed,
             text_cfg,
             img_cfg,
@@ -173,13 +164,13 @@ with gr.Blocks(css=css) as demo:
             positive_prompt,
             negative_prompt,
         ],
-        outputs=[result_1, result_2],
+        outputs=[result_1],
     )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=show_result,
-        outputs=[result_1, result_2],
+        outputs=[result_1],
     )