Update app.py

app.py

@@ -161,7 +161,8 @@ class ForgeLoader4Bit(torch.nn.Module):
                 data=state_dict[prefix + 'weight'],
                 quantized_stats=quant_state_dict,
                 requires_grad=False,
-                device=self.dummy.device,
+                # device=self.dummy.device,
+                device=torch.device('cuda'),
                 module=self
             )
             self.quant_state = self.weight.quant_state
@@ -717,12 +718,10 @@ def get_image(image) -> torch.Tensor | None:
 # ---------------- Demo ----------------
 
 
-from pathlib import Path
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
-if not Path("flux1-dev-bnb-nf4.safetensors").exists():
-    torch.hub.download_url_to_file("https://huggingface.co/lllyasviel/flux1-dev-bnb-nf4/resolve/main/flux1-dev-bnb-nf4.safetensors", "flux1-dev-bnb-nf4.safetensors")
-    
-sd = load_file("flux1-dev-bnb-nf4.safetensors")
+sd = load_file(hf_hub_download(repo_id="lllyasviel/flux1-dev-bnb-nf4", filename="flux1-dev-bnb-nf4.safetensors"))
 sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items() if "model.diffusion_model" in k}
 model = Flux().to(dtype=torch.float16, device="cuda")
 result = model.load_state_dict(sd)
@@ -753,7 +752,8 @@ def generate_image(
             init_image = init_image[..., : 16 * (h // 16), : 16 * (w // 16)]
             height = init_image.shape[-2]
             width = init_image.shape[-1]
-        init_image = ae.encode(init_image.to(torch_device))
+        init_image = ae.encode(init_image.to(torch_device)).latent_dist.sample()
+        init_image =  (init_image - ae.config.shift_factor) * ae.config.scaling_factor
 
     generator = torch.Generator(device=device).manual_seed(seed)
     x = torch.randn(1, 16, 2 * math.ceil(height / 16), 2 * math.ceil(width / 16), device=device, dtype=torch.bfloat16, generator=generator)