22

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import numpy as np
 import torch, random, json, spaces, time
+from ulid import ULID
 from diffsynth.pipelines.qwen_image import (
     QwenImagePipeline, ModelConfig,
     QwenImageUnit_Image2LoRAEncode, QwenImageUnit_Image2LoRADecode
@@ -14,6 +15,7 @@ from utils import repo_utils, image_utils, prompt_utils
 # repo_utils.clone_repo_if_not_exists("git clone https://huggingface.co/DiffSynth-Studio/General-Image-Encoders", "app/repos")
 # repo_utils.clone_repo_if_not_exists("https://huggingface.co/apple/starflow", "app/models")
 
+URL_PUBLIC = "https://huggingface.co/spaces/AiSudo/Qwen-Image-to-LoRA/blob/main"
 DTYPE = torch.bfloat16
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -57,35 +59,14 @@ pipe = QwenImagePipeline.from_pretrained(
 )
 
 
-
-# pipe = ZImageControlPipeline(
-#     vae=vae,
-#     tokenizer=tokenizer,
-#     text_encoder=text_encoder,
-#     transformer=transformer,
-#     scheduler=scheduler,
-# )
-# pipe.to("cuda", DTYPE)
-
-
-# def prepare(prompt, is_polish_prompt):
-#     if not is_polish_prompt: return prompt, False
-#     polished_prompt = prompt_utils.polish_prompt(prompt)
-#     return polished_prompt, True
-
-
 @spaces.GPU
-def inference(
-    prompt,
-    negative_prompt,
-    seed=42,
-    randomize_seed=True,
-    guidance_scale=1.5,
-    num_inference_steps=8,
+def generate_lora(
+    images,
     progress=gr.Progress(track_tqdm=True),
 ):
-    timestamp = time.time()
-    print(f"timestamp: {timestamp}")
+
+    ulid = str(ULID()).lower()
+    print(f"ulid: {ulid}")
 
     # Load images 
     images = [
@@ -102,40 +83,25 @@ def inference(
         embs = QwenImageUnit_Image2LoRAEncode().process(pipe, image2lora_images=images)
         lora = QwenImageUnit_Image2LoRADecode().process(pipe, **embs)["lora"]
 
-    save_file(lora, "model_style.safetensors")
-
-    return True
+    lora_path = f"loras/{ulid}.safetensors"
+    lora_url = f"{URL_PUBLIC}/{lora_path}"
 
-#     # process image
-#     print("DEBUG: process image")
-#     if input_image is None:
-#         print("Error: input_image is empty.")
-#         return None
+    save_file(lora, lora_path)
 
+    return True
 
-#     print("DEBUG: control_image_torch")
-#     orig_width, orig_height = input_image.size
-#     control_image, width, height = image_utils.rescale_image(input_image, image_scale, 16, 2048)
-#     control_image_torch = image_utils.get_image_latent(control_image, sample_size=[height, width])[:, :, 0]
+@spaces.GPU
+def generate_image(
+    prompt,
+    negative_prompt,
+    seed=42,
+    randomize_seed=True,
+    guidance_scale=1.5,
+    num_inference_steps=8,
+    progress=gr.Progress(track_tqdm=True),
+):
     
-#     # generation
-#     if randomize_seed: seed = random.randint(0, MAX_SEED)
-#     generator = torch.Generator().manual_seed(seed)
-
-#     output_image = pipe(
-#         prompt=prompt,
-#         negative_prompt = negative_prompt,
-#         width=width,
-#         height=height,
-#         generator=generator,
-#         guidance_scale=guidance_scale,
-#         control_image=control_image_torch,
-#         num_inference_steps=num_inference_steps,
-#         control_context_scale=control_context_scale,
-#     ).images[0]
-
-#     output_image = output_image.resize((orig_width * image_scale, orig_height * image_scale))
-#     return output_image, seed
+    return True
 
 
 def read_file(path: str) -> str:
@@ -151,15 +117,17 @@ css = """
 }
 """
 
-with open('examples/0_examples.json', 'r') as file: examples = json.load(file)
 
+
+with open('examples/0_examples.json', 'r') as file: examples = json.load(file)
 with gr.Blocks() as demo:
+    has_lora = True
     with gr.Column(elem_id="col-container"):
         with gr.Column():
             gr.HTML(read_file("static/header.html"))
         with gr.Row():
             with gr.Column():
-                gallery = gr.Gallery(
+                input_images = gr.Gallery(
                     label="Generated images", 
                     show_label=False, 
                     elem_id="gallery", 
@@ -167,6 +135,15 @@ with gr.Blocks() as demo:
                     object_fit="cover", 
                     height=300)
                 
+                lora_button = gr.Button("Generate LoRA", variant="primary")
+
+            with gr.Column():
+                lora_path = gr.Textbox(label="Generated LoRA path",lines=2, interactive=False)
+                lora_download = gr.DownloadButton(label=f"Download LoRA", visible=has_lora)
+        
+        with gr.Row(visible=has_lora):
+            gr.Markdown("Your LoRA is ready! Now, try generating some images.")
+            with gr.Column():
                 prompt = gr.Textbox(
                     label="Prompt",
                     show_label=False,
@@ -175,13 +152,10 @@ with gr.Blocks() as demo:
                     value="a man in a fishing boat. high quality, detailed"
                     # container=False,
                 )
-                # is_polish_prompt = gr.Checkbox(label="Polish prompt", value=True)
-                # control_mode = gr.Radio(
-                #     choices=["Canny", "Depth", "HED", "MLSD", "Pose"],
-                #     value="Canny",
-                #     label="Control Mode"
-                # )
-                run_button = gr.Button("Generate", variant="primary")
+                
+                imagen_button = gr.Button("Generate Image", variant="primary")
+
+                
                 with gr.Accordion("Advanced Settings", open=False):
                     
                     negative_prompt = gr.Textbox(
@@ -227,26 +201,16 @@ with gr.Blocks() as demo:
 
             with gr.Column():
                 output_image = gr.Image(label="Generated image", show_label=False)
-                # polished_prompt = gr.Textbox(label="Polished prompt", interactive=False)
-
-                # with gr.Accordion("Preprocessor output", open=False):
-                #     control_image = gr.Image(label="Control image", show_label=False)
-
                     
         # gr.Examples(examples=examples, inputs=[input_image])
         gr.Markdown(read_file("static/footer.md"))
 
-    run_button.click(
-        fn=inference,
+    lora_button.click(
+        fn=generate_lora,
         inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            guidance_scale,
-            num_inference_steps,
+            input_images
         ],
-        outputs=[output_image, seed],
+        outputs=[lora_path, lora_download, has_lora],
     )