Update app.py

app.py

@@ -9,44 +9,15 @@ from safetensors.torch import load_file
 from omegaconf import OmegaConf
 
 from image_datasets.dataset import image_resize
-def tensor_to_pil_image(in_image):
-    tensor = in_image.squeeze(0)
-    tensor = (tensor + 1) / 2
-    tensor = tensor * 255
-    numpy_array = tensor.permute(1, 2, 0).byte().numpy()
-    pil_image = Image.fromarray(numpy_array)
-    return pil_image
-# from src.flux.xflux_pipeline import XFluxSampler
+
 args = OmegaConf.load("inference_configs/inference.yaml")
-# 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)
-#test push
 @spaces.GPU
 def generate(image: Image.Image, edit_prompt: str):
     from src.flux.xflux_pipeline import XFluxSampler
 
     
-    
-
-    # 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)
 
     sampler = XFluxSampler(
         device = device,
@@ -56,42 +27,7 @@ def generate(image: Image.Image, edit_prompt: str):
         image_encoder=None,
         improj=None
     )
-    # global sampler
-    # device = torch.device("cuda")
-    # dtype = torch.bfloat16
-
-    # if sampler is None:
-        # 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)
-
-        # 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()
-
-        # sampler = XFluxSampler(
-        #     clip=clip,
-        #     t5=t5,
-        #     ae=vae,
-        #     model=dit,
-        #     device=device,
-        #     ip_loaded=False,
-        #     spatial_condition=False,
-        #     clip_image_processor=None,
-        #     image_encoder=None,
-        #     improj=None
-        # )
-
+    
     img = image_resize(image, 512)
     w, h = img.size
     img = img.resize(((w // 32) * 32, (h // 32) * 32))