new

app.py

@@ -146,6 +146,7 @@ def text_to_detailed(prompt, seed=None):
     # print(f"Before text_to_detailed: {torch.cuda.memory_allocated() / 1024**3} GB")
     return k3d_wrapper.get_detailed_prompt(prompt, seed)
 
+@spaces.GPU
 def text_to_image(prompt, seed=None, strength=1.0,lora_scale=1.0, num_inference_steps=30, redux_hparam=None, init_image=None, **kwargs):
     # print(f"Before text_to_image: {torch.cuda.memory_allocated() / 1024**3} GB")
     k3d_wrapper.renew_uuid()

pipeline/kiss3d_wrapper.py

@@ -70,22 +70,22 @@ def init_wrapper_from_config(config_path):
     flux_lora_pth = hf_hub_download(repo_id="LTT/Kiss3DGen", filename="rgb_normal_large.safetensors", repo_type="model", token=access_token)
     flux_redux_pth = config_['flux'].get('redux', None)
     
-    # if flux_base_model_pth.endswith('safetensors'):
-    #     flux_pipe = FluxImg2ImgPipeline.from_single_file(flux_base_model_pth, torch_dtype=dtype_[flux_dtype], token=access_token)
-    # else:
-    #     flux_pipe = FluxImg2ImgPipeline.from_pretrained(flux_base_model_pth, torch_dtype=dtype_[flux_dtype], token=access_token)
+    if flux_base_model_pth.endswith('safetensors'):
+        flux_pipe = FluxImg2ImgPipeline.from_single_file(flux_base_model_pth, torch_dtype=dtype_[flux_dtype], token=access_token)
+    else:
+        flux_pipe = FluxImg2ImgPipeline.from_pretrained(flux_base_model_pth, torch_dtype=dtype_[flux_dtype], token=access_token)
     
     # load flux model and controlnet
     if flux_controlnet_pth is not None and False:
         flux_controlnet = FluxControlNetModel.from_pretrained(flux_controlnet_pth, torch_dtype=torch.bfloat16)
         flux_pipe = convert_flux_pipeline(flux_pipe, FluxControlNetImg2ImgPipeline, controlnet=[flux_controlnet])
 
-    # flux_pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(flux_pipe.scheduler.config)
+    flux_pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(flux_pipe.scheduler.config)
         
     # load lora weights
-    # flux_pipe.load_lora_weights(flux_lora_pth)
-    # flux_pipe.to(device=flux_device)
-    flux_pipe = None
+    flux_pipe.load_lora_weights(flux_lora_pth)
+    flux_pipe.to(device=flux_device)
+    # flux_pipe = None
 
     # load redux model
     flux_redux_pipe = None
@@ -465,7 +465,6 @@ class kiss3d_wrapper(object):
     
         return preprocessed
 
-    @spaces.GPU
     def generate_3d_bundle_image_text(self, 
                                       prompt,
                                       image=None,