Update app.py

app.py

@@ -12,16 +12,14 @@ from image_datasets.dataset import image_resize
 from src.flux.util import load_ae, load_clip, load_flow_model2, load_t5, tensor_to_pil_image
 from src.flux.xflux_pipeline import XFluxSampler
 args = OmegaConf.load("inference_configs/inference.yaml")
-is_schnell = args.model_name == "flux-schnell"
-'/home/user/app/assets/0_camera_zoom/20486354.png'
-'/home/user/app/assets/0_camera_zoom/20486354.png'
+# is_schnell = args.model_name == "flux-schnell"
 # sampler = None
-device = torch.device("cuda")
-dtype = torch.bfloat16
-dit = load_flow_model2(args.model_name, device="cpu").to(device, dtype=dtype)
-vae = load_ae(args.model_name, device="cpu").to(device, dtype=dtype)
-t5 = load_t5(device="cpu", max_length=256 if is_schnell else 512).to(device, dtype=dtype)
-clip = load_clip("cpu").to(device, dtype=dtype)
+# device = torch.device("cuda")
+# dtype = torch.bfloat16
+# dit = load_flow_model2(args.model_name, device="cpu").to(device, dtype=dtype)
+# vae = load_ae(args.model_name, device="cpu").to(device, dtype=dtype)
+# t5 = load_t5(device="cpu", max_length=256 if is_schnell else 512).to(device, dtype=dtype)
+# clip = load_clip("cpu").to(device, dtype=dtype)
 #test push
 @spaces.GPU
 def generate(image: Image.Image, edit_prompt: str):
@@ -29,26 +27,21 @@ def generate(image: Image.Image, edit_prompt: str):
     
     
 
-    vae.requires_grad_(False)
-    t5.requires_grad_(False)
-    clip.requires_grad_(False)
+    # vae.requires_grad_(False)
+    # t5.requires_grad_(False)
+    # clip.requires_grad_(False)
 
-    model_path = hf_hub_download(
-        repo_id="Boese0601/ByteMorpher",
-        filename="dit.safetensors",
-        use_auth_token=os.getenv("HF_TOKEN")
-    )
-    state_dict = load_file(model_path)
-    dit.load_state_dict(state_dict)
-    dit.eval()
-    dit.to(device, dtype=dtype)
+    # model_path = hf_hub_download(
+    #     repo_id="Boese0601/ByteMorpher",
+    #     filename="dit.safetensors",
+    #     use_auth_token=os.getenv("HF_TOKEN")
+    # )
+    # state_dict = load_file(model_path)
+    # dit.load_state_dict(state_dict)
+    # dit.eval()
+    # dit.to(device, dtype=dtype)
 
     sampler = XFluxSampler(
-        clip=clip,
-        t5=t5,
-        ae=vae,
-        model=dit,
-        device=device,
         ip_loaded=False,
         spatial_condition=False,
         clip_image_processor=None,