Update app.py

app.py

@@ -1,14 +1,13 @@
 import os
-if os.environ.get("SPACES_ZERO_GPU") is not None:
-    import spaces
-else:
-    class spaces:
-        @staticmethod
-        def GPU(func):
-            def wrapper(*args, **kwargs):
-                return func(*args, **kwargs)
-            return wrapper
-
+#if os.environ.get("SPACES_ZERO_GPU") is not None:
+import spaces
+#else:
+#    class spaces:
+#        @staticmethod
+#        def GPU(func):
+#            def wrapper(*args, **kwargs):
+#                return func(*args, **kwargs)
+#            return wrapper
 import gradio as gr
 import json
 import logging
@@ -43,11 +42,13 @@ from diffusers.models.autoencoders import AutoencoderKL
 from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
 import gc
 import warnings
-model_path = snapshot_download(repo_id="Kijai/OpenFLUX-comfy")
-cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
-os.environ["TRANSFORMERS_CACHE"] = cache_path
-os.environ["HF_HUB_CACHE"] = cache_path
-os.environ["HF_HOME"] = cache_path
+#model_path = snapshot_download(repo_id="Kijai/OpenFLUX-comfy")
+model_path  = snapshot_download(repo_id="nyanko7/flux-dev-de-distill")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+#cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
+#os.environ["TRANSFORMERS_CACHE"] = cache_path
+#os.environ["HF_HUB_CACHE"] = cache_path
+#os.environ["HF_HOME"] = cache_path
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -690,7 +691,7 @@ class EmptyInitWrapper(torch.overrides.TorchFunctionMode):
 with EmptyInitWrapper():
     model = Flux().to(dtype=torch.bfloat16, device="cuda")
     
-    sd = load_file(f"{model_path}/OpenFlux-fp8_e4m3fn.safetensors")
+    sd = load_file(f"{model_path}/consolidated_s6700.safetensors")
     sd = {k.replace("model.", ""): v for k, v in sd.items()}
     result = model.load_state_dict(sd)