bugfix

app.py

@@ -18,7 +18,8 @@ import time
 import boto3
 from io import BytesIO
 from datetime import datetime
-from diffusers.utils import load_image
+from diffusers.utils import load_image, make_image_grid
+
 import json
 from preprocessor import Preprocessor
 from diffusers.pipelines.flux.pipeline_flux_controlnet_inpaint import FluxControlNetInpaintPipeline
@@ -32,17 +33,14 @@ MAX_SEED = np.iinfo(np.int32).max
 IMAGE_SIZE = 1024
 
 # init
-dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
 controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union-alpha'
 controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
 
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
-pipe = FluxControlNetInpaintPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=dtype, vae=taef1).to(device)
 
+pipe = FluxControlNetInpaintPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16).to(device)
 
 
 control_mode_ids = {
@@ -176,7 +174,7 @@ def run_flux(
     generator = torch.Generator().manual_seed(seed_slicer)
 
     with calculateDuration("run pipe"):
-        genearte_image = pipe(
+        generated_image = pipe(
             prompt=prompt,
             image=image,
             mask_image=mask,
@@ -191,7 +189,7 @@ def run_flux(
             joint_attention_kwargs={"scale": lora_scale}
         ).images[0]
     
-    return genearte_image
+    return generated_image
 
 @spaces.GPU(duration=120)
 def process(
@@ -306,19 +304,30 @@ def process(
             resolution_wh=(width, height),
             progress=progress
         )
-    except:
+    except Exception as e:
+        result["status"] = "faield"
         result["message"] = "generate image failed"
-        return None, json.dumps(result)
+        print(e)
+        generated_image = None
 
     print("run flux finish")
-    if upload_to_r2:
-        with calculateDuration("upload image"):
-            url = upload_image_to_r2(generated_image, account_id, access_key, secret_key, bucket)
-            result = {"status": "success", "message": "upload image success", "url": url}
-    else:
-        result = {"status": "success", "message": "Image generated but not uploaded"}
+    if generated_image:
+        if upload_to_r2:
+            with calculateDuration("upload image"):
+                url = upload_image_to_r2(generated_image, account_id, access_key, secret_key, bucket)
+                result = {"status": "success", "message": "upload image success", "url": url}
+        else:
+            result = {"status": "success", "message": "Image generated but not uploaded"}
+
 
-    return generated_image, json.dumps(result)
+    final_images = []
+    final_images.append(image)
+    final_images.append(mask)
+    final_images.append(control_image)
+    if generated_image:
+        final_images.append(generated_image)
+        
+    return final_images, json.dumps(result)
 
 
 with gr.Blocks() as demo:
@@ -449,7 +458,7 @@ with gr.Blocks() as demo:
                     secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
                 
         with gr.Column():
-            generated_image = gr.Image(label="Result", show_label=False)
+            generated_images = gr.Gallery(label="Result", show_label=True)
             output_json_component = gr.Code(label="JSON Result", language="json")
 
     submit_button_component.click(
@@ -475,7 +484,7 @@ with gr.Blocks() as demo:
             bucket
         ],
         outputs=[
-            generated_image,
+            generated_images,
             output_json_component
         ]
     )