Update app.py

app.py

@@ -59,10 +59,11 @@ class calculateDuration:
 @spaces.GPU(duration=120)
 @torch.inference_mode()
 def generate_image(prompt, steps, seed, cfg_scale, width, height, progress):
-    pipe.to(device)
-    generator = torch.Generator(device=device).manual_seed(seed)
+    
     with calculateDuration("Generating image"):
         # Generate image
+        pipe.to(device)
+        generator = torch.Generator(device=device).manual_seed(seed)
         print(prompt, steps, seed, cfg_scale, width, height)
         generated_image = pipe(
             prompt=prompt,
@@ -100,6 +101,7 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     return image_file
 
 def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
+    print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height)
     
     # Load LoRA weights
     if lora_strings_json:
@@ -131,9 +133,11 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
             seed = random.randint(0, MAX_SEED)
     
     # Generate image
+    error_message = ""
     try:
         final_image = generate_image(prompt, steps, seed, cfg_scale, width, height, progress)
-    except:
+    except Exception as e:
+        error_message =  str(e)
         final_image = None
         
     if final_image:
@@ -144,7 +148,7 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
         else:
             result = {"status": "success", "message": "Image generated but not uploaded"}
     else:
-        result = {"status": "failed", "message": "Image generate failed"}
+        result = {"status": "failed", "message": error_message}
         
     progress(100, "Completed!")
 
@@ -165,8 +169,8 @@ with gr.Blocks(css=css) as demo:
         
         with gr.Column():
 
-            prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=2)
-            lora_strings_json = gr.Text(label="LoRA Strings (JSON List)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1"}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2"}]', lines=5)
+            prompt = gr.Text(label="Prompt", placeholder="Enter prompt", lines=10)
+            lora_strings_json = gr.Text(label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1", "adapter_weight": 1}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2", "adapter_weight": 1}]', lines=5)
             
             run_button = gr.Button("Run", scale=0)