main

gradio_app_asy.py

@@ -1,4 +1,4 @@
-# import spaces
+import spaces
 import gradio as gr
 import torch
 import numpy as np
@@ -25,48 +25,8 @@ logging.basicConfig(level=logging.DEBUG)
 
 accelerator = Accelerator(mixed_precision='bf16', device_placement=True)
 
-# hf_token = os.getenv("HF_TOKEN")
-# login(token=hf_token)
-
-# # Model paths dynamically retrieved using selected model
-# model_paths = {
-#     'Wood Sculpture': {
-#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
-#         'BASE_FILE': "flux_merge_lora/flux_merge_4f_wood-fp16.safetensors",
-#         'LORA_REPO': "showlab/makeanything",
-#         'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors",
-#         "Frame": 4
-#     },
-#     'LEGO': {
-#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
-#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_lego-fp16.safetensors",
-#         'LORA_REPO': "showlab/makeanything",
-#         'LORA_FILE': "recraft/recraft_9f_lego.safetensors",
-#         "Frame": 9
-#     },
-#     'Sketch': {
-#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
-#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_portrait-fp16.safetensors",
-#         'LORA_REPO': "showlab/makeanything",
-#         'LORA_FILE': "recraft/recraft_9f_sketch.safetensors",
-#         "Frame": 9
-#     },
-#     'Portrait': {
-#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
-#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_sketch-fp16.safetensors",
-#         'LORA_REPO': "showlab/makeanything",
-#         'LORA_FILE': "recraft/recraft_9f_portrait.safetensors",
-#         "Frame": 9
-#     }
-# }
-
-# # Common paths
-# clip_repo_id = "comfyanonymous/flux_text_encoders"
-# t5xxl_file = "t5xxl_fp16.safetensors"
-# clip_l_file = "clip_l.safetensors"
-# ae_repo_id = "black-forest-labs/FLUX.1-dev"
-# ae_file = "ae.safetensors"
-
+hf_token = os.getenv("HF_TOKEN")
+login(token=hf_token)
 domain_index = {
     'LEGO': 1, 'Cook': 2, 'Painting': 3, 'Icon': 4, 'Landscape illustration': 5, 
     'Portrait': 6, 'Transformer': 7, 'Sand art': 8, 'Illustration': 9,  'Sketch': 10,
@@ -76,15 +36,19 @@ domain_index = {
 }
 
 lora_paths = {
-    "9 frame": "/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_9f_general.safetensors",
-    "4 frame": "/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_4f_general.safetensors"
+    "9 frame": "asymmetric_lora/asymmetric_lora_9f_general.safetensors",
+    "4 frame": "asymmetric_lora/asymmetric_lora_4f_general.safetensors"
 }
-BASE_FLUX_CHECKPOINT = "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/unet/flux1-dev-fp8.safetensors"
-# LORA_WEIGHTS_PATH="/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_9f_general.safetensors"
-CLIP_L_PATH = "/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/clip_l.safetensors"
-T5XXL_PATH = "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/clip/t5xxl_fp8_e4m3fn.safetensors"
-AE_PATH = "/tiamat-vePFS/share_data/storage/huggingface/models/black-forest-labs/FLUX.1-dev/ae.safetensors"
 
+# Common paths
+flux_repo_id="Kijai/flux-fp8"
+flux_file="flux1-dev-fp8.safetensors"
+lora_repo_id="showlab/makeanything"
+clip_repo_id = "comfyanonymous/flux_text_encoders"
+t5xxl_file = "t5xxl_fp16.safetensors"
+clip_l_file = "clip_l.safetensors"
+ae_repo_id = "black-forest-labs/FLUX.1-dev"
+ae_file = "ae.safetensors"
 
 # Model placeholders
 model = None
@@ -100,22 +64,25 @@ def download_file(repo_id, file_name):
 # Load model function with dynamic paths based on the selected model
 def load_target_model(frame, domain):
     global model, clip_l, t5xxl, ae, lora_model
+    BASE_FLUX_CHECKPOINT=download_file(flux_repo_id, flux_file)
+    CLIP_L_PATH = download_file(clip_repo_id, clip_l_file)
+    T5XXL_PATH = download_file(clip_repo_id, t5xxl_file)
+    AE_PATH = download_file(ae_repo_id, ae_file)
+    LORA_WEIGHTS_PATH = download_file(lora_repo_id, lora_paths[frame])
 
     logger.info("Loading models...")
     # try:
     if model is None is None or clip_l is None or t5xxl is None or ae is None:
+        _, model = flux_utils.load_flow_model(
+            BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
+        )
         clip_l = flux_utils.load_clip_l(CLIP_L_PATH, torch.bfloat16, "cpu", disable_mmap=False)
         clip_l.eval()
         t5xxl = flux_utils.load_t5xxl(T5XXL_PATH, torch.bfloat16, "cpu", disable_mmap=False)
         t5xxl.eval()
         ae = flux_utils.load_ae(AE_PATH, torch.bfloat16, "cpu", disable_mmap=False)
         logger.info("Models loaded successfully.")
-    # Load models
-    _, model = flux_utils.load_flow_model(
-        BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
-    )
     # Load LoRA weights
-    LORA_WEIGHTS_PATH = lora_paths[frame]
     multiplier = 1.0
     weights_sd = load_file(LORA_WEIGHTS_PATH)
     lora_ups_num = 10 if frame=="9 frame" else 21
@@ -131,12 +98,8 @@ def load_target_model(frame, domain):
     logger.info("Models loaded successfully.")
     return "Models loaded successfully. Using Frame: {}, Damain: {}".format(frame, domain)
 
-    # except Exception as e:
-    #     logger.error(f"Error loading models: {e}")
-    #     return f"Error loading models: {e}"
-
 # The function to generate image from a prompt and conditional image
-# @spaces.GPU(duration=180)
+@spaces.GPU(duration=180)
 def infer(prompt, frame, seed=0):
     global model, clip_l, t5xxl, ae, lora_model
     if model is None or lora_model is None or clip_l is None or t5xxl is None or ae is None:
@@ -257,19 +220,14 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    frame_selector = gr.Radio(choices=["4 frame", "9 frame"], label="Select Floor")
-                with gr.Column(scale=2):
-                    domain_selector = gr.Dropdown(choices=[], label="Select Domains")
-                   
-            # Load Model Button
-            load_button = gr.Button("Load Model")
-
+            frame_selector = gr.Radio(choices=["4 frame", "9 frame"], label="Select Model")
         with gr.Column(scale=1):
-            # Status message box
-            status_box = gr.Textbox(label="Status", placeholder="Model loading status", interactive=False, value="Model not loaded", lines=3)
-    
+                load_button = gr.Button("Load Model")
+        with gr.Column(scale=2):
+            status_box = gr.Textbox(label="Status", placeholder="Model loading status", interactive=False, value="Model not loaded", lines=2)
+        with gr.Column(scale=2):
+            domain_selector = gr.Dropdown(choices=[], label="Select Domains")
+
     with gr.Row():
         with gr.Column(scale=1):
             # Input for the prompt
@@ -326,4 +284,4 @@ with gr.Blocks() as demo:
     # )
 
     # Launch the Gradio app
-    demo.launch(server_port=8289, server_name="0.0.0.0", share=True)
+    demo.launch()