Update app.py

app.py

@@ -20,18 +20,18 @@ from io import BytesIO
 import re
 import json
 
-# 登录 Hugging Face Hub
+# Login Hugging Face Hub
 HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 import diffusers
 print(diffusers.__version__)
 
-# 初始化
-dtype = torch.float16  # 您可以根据需要调整数据类型
+# init
+dtype = torch.float16  # use float16 for fast generate
 device = "cuda" if torch.cuda.is_available() else "cpu"
-base_model = "black-forest-labs/FLUX.1-dev"  # 替换为您的模型
+base_model = "black-forest-labs/FLUX.1-dev"
 
-# 加载管道
+# load pipe
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
 
 MAX_SEED = 2**32 - 1
@@ -56,10 +56,7 @@ class calculateDuration:
         else:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
         
-        print(f"Activity: {self.activity_name}, End time: {self.start_time_formatted}")
-
-# 生成图像的函数
-@spaces.GPU
+@spaces.GPU(duration=120)
 @torch.inference_mode()
 def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
     pipe.to(device)
@@ -121,10 +118,11 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
                     adapter_name = lora_info.get("adapter_name")
                     adapter_weight = lora_info.get("adapter_weight")
                     if lora_repo and weights and adapter_name:
-                        # 加载 LoRA 权重
+                        # load lora
                         pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
                         adapter_names.append(adapter_name)
                         adapter_weights.append(adapter_weight)
+            # set lora weights
             pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
     
     # Set random seed for reproducibility
@@ -133,8 +131,11 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
             seed = random.randint(0, MAX_SEED)
     
     # Generate image
-    final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
-    
+    try:
+        final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
+    except:
+        final_image = None
+        
     if final_image:
         if upload_to_r2:
             with calculateDuration("Upload image"):
@@ -142,12 +143,15 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
                 result = {"status": "success", "message": "upload image success", "url": url}
         else:
             result = {"status": "success", "message": "Image generated but not uploaded"}
-
+    else:
+        result = {"status": "failed", "message": "Image generate failed"}
+        
     progress(100, "Completed!")
 
     return final_image, seed, json.dumps(result)
 
-# Gradio 界面
+# Gradio interface
+
 css="""
 #col-container {
     margin: 0 auto;
@@ -156,7 +160,7 @@ css="""
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("Flux with LoRA")
+    gr.Markdown("flux-dev-multi-lora")
     with gr.Row():
         
         with gr.Column():