Update app.py

app.py

@@ -202,19 +202,23 @@ def remove_lora_2(selected_indices, loras_state):
     return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2
 
 def randomize_loras(selected_indices, loras_state):
-    if len(loras_state) < 2:
-        raise gr.Error("Not enough LoRAs to randomize.")
-    selected_indices = random.sample(range(len(loras_state)), 2)
-    lora1 = loras_state[selected_indices[0]]
-    lora2 = loras_state[selected_indices[1]]
-    selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
-    selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
-    lora_scale_1 = 1.15
-    lora_scale_2 = 1.15
-    lora_image_1 = lora1['image']
-    lora_image_2 = lora2['image']
-    random_prompt = random.choice(prompt_values)
-    return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
+    try:
+        if len(loras_state) < 2:
+            raise gr.Error("Not enough LoRAs to randomize.")
+        selected_indices = random.sample(range(len(loras_state)), 2)
+        lora1 = loras_state[selected_indices[0]]
+        lora2 = loras_state[selected_indices[1]]
+        selected_info_1 = f"### LoRA 1 Selected: [{lora1['title']}](https://huggingface.co/{lora1['repo']}) ✨"
+        selected_info_2 = f"### LoRA 2 Selected: [{lora2['title']}](https://huggingface.co/{lora2['repo']}) ✨"
+        lora_scale_1 = 1.15
+        lora_scale_2 = 1.15
+        lora_image_1 = lora1['image']
+        lora_image_2 = lora2['image']
+        random_prompt = random.choice(prompt_values)
+        return selected_info_1, selected_info_2, selected_indices, lora_scale_1, lora_scale_2, lora_image_1, lora_image_2, random_prompt
+    except Exception as e:
+        print(f"Error in randomize_loras: {str(e)}")
+        return "Error", "Error", [], 1.15, 1.15, None, None, ""
 
 def add_custom_lora(custom_lora, selected_indices, current_loras):
     if custom_lora:
@@ -566,42 +570,52 @@ def infer_upscale(
     controlnet_conditioning_scale,
     progress=gr.Progress(track_tqdm=True),
 ):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    true_input_image = input_image
-    input_image, w_original, h_original, was_resized = process_input(
-        input_image, upscale_factor
-    )
+    try:
+        if input_image is None:
+            raise ValueError("No input image provided")
+        
+        if randomize_seed:
+            seed = random.randint(0, MAX_SEED)
+        true_input_image = input_image
+        input_image, w_original, h_original, was_resized = process_input(
+            input_image, upscale_factor
+        )
 
-    # rescale with upscale factor
-    w, h = input_image.size
-    control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
-
-    generator = torch.Generator().manual_seed(seed)
-
-    gr.Info("Upscaling image...")
-    image = pipe_upscale(
-        prompt="",
-        control_image=control_image,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=3.5,
-        height=control_image.size[1],
-        width=control_image.size[0],
-        generator=generator,
-    ).images[0]
+        # rescale with upscale factor
+        w, h = input_image.size
+        control_image = input_image.resize((w * upscale_factor, h * upscale_factor))
 
-    if was_resized:
-        gr.Info(
-            f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
-        )
+        generator = torch.Generator(device=device).manual_seed(seed)
+
+        gr.Info("Upscaling image...")
+        # 모든 텐서를 동일한 디바이스로 이동
+        pipe_upscale.to(device)
+        control_image = control_image.to(device)
+        
+        image = pipe_upscale(
+            prompt="",
+            control_image=control_image,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=3.5,
+            height=control_image.size[1],
+            width=control_image.size[0],
+            generator=generator,
+        ).images[0]
 
-    # resize to target desired size
-    image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
-    image.save("output.jpg")
-    # convert to numpy
-    return [true_input_image, image, seed]
+        if was_resized:
+            gr.Info(
+                f"Resizing output image to targeted {w_original * upscale_factor}x{h_original * upscale_factor} size."
+            )
 
+        # resize to target desired size
+        image = image.resize((w_original * upscale_factor, h_original * upscale_factor))
+        image.save("output.jpg")
+        # convert to numpy
+        return [true_input_image, image, seed]
+    except Exception as e:
+        print(f"Error in infer_upscale: {str(e)}")
+        return [None, None, seed]
 
 
 with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css, delete_cache=(60, 3600)) as app: