Update app.py

app.py

@@ -50,22 +50,7 @@ MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 os.makedirs(TMP_DIR, exist_ok=True)
 
-controlnet = ControlNetModel.from_pretrained(
-    "xinsir/controlnet-scribble-sdxl-1.0",
-    torch_dtype=torch.float16
-)
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-
-pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-    "sd-community/sdxl-flash",
-    controlnet=controlnet,
-    vae=vae,
-    torch_dtype=torch.float16,
-    # scheduler=eulera_scheduler,
-)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-pipe.to('cuda')
 
 def start_session(req: gr.Request):
     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
@@ -78,7 +63,7 @@ def end_session(req: gr.Request):
 
 
 def preprocess_image(image: Image.Image, 
-                     prompt: str,
+                    prompt: str,
                     negative_prompt: str = "",
                     num_steps: int = 25,
                     guidance_scale: float = 5,
@@ -92,21 +77,23 @@ def preprocess_image(image: Image.Image,
     Returns:
         Image.Image: The preprocessed image.
     """
-    width, height  = image['composite'].size
-    ratio = np.sqrt(1024. * 1024. / (width * height))
-    new_width, new_height = int(width * ratio), int(height * ratio)
-    image = image['composite'].resize((new_width, new_height))
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        num_inference_steps=num_steps,
-        controlnet_conditioning_scale=controlnet_conditioning_scale,
-        guidance_scale=guidance_scale,
-        width=new_width,
-        height=new_height).images[0]
-    print(type(image))
+    if prompt is not None:
+        width, height  = image['composite'].size
+        ratio = np.sqrt(1024. * 1024. / (width * height))
+        new_width, new_height = int(width * ratio), int(height * ratio)
+        image = image['composite'].resize((new_width, new_height))
+    
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            num_inference_steps=num_steps,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guidance_scale=guidance_scale,
+            width=new_width,
+            height=new_height).images[0]
+        print(type(image))
+        
     processed_image = pipeline.preprocess_image(image)
     return processed_image
 
@@ -481,8 +468,26 @@ with gr.Blocks(delete_cache=(600, 600), js=js_func) as demo:
 if __name__ == "__main__":
     pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
     pipeline.cuda()
-    try:
-        pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
-    except:
-        pass
+    
+    #scribble controlnet
+    controlnet = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-scribble-sdxl-1.0",
+    torch_dtype=torch.float16
+)
+    vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+    
+    pipe_control = StableDiffusionXLControlNetPipeline.from_pretrained(
+        "sd-community/sdxl-flash",
+        controlnet=controlnet,
+        vae=vae,
+        torch_dtype=torch.float16,
+        # scheduler=eulera_scheduler,
+    )
+    pipe_control.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_control.scheduler.config)
+    pipe_control.cuda()
+    
+    # try:
+    #     pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))    # Preload rembg
+    # except:
+    #     pass
     demo.launch()