Update app.py

app.py

@@ -2,7 +2,6 @@ import logging
 import random
 import warnings
 import os
-os.environ["GIT_TERMINAL_PROMPT"] = "0"
 import gradio as gr
 import numpy as np
 import spaces
@@ -15,6 +14,8 @@ import subprocess
 from huggingface_hub import hf_hub_download
 import tempfile
 
+os.environ["GIT_TERMINAL_PROMPT"] = "0"
+
 # Setup ComfyUI and custom nodes
 if not os.path.exists("ComfyUI"):
     subprocess.run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI"])
@@ -22,10 +23,10 @@ if not os.path.exists("ComfyUI"):
 custom_nodes_dir = os.path.join("ComfyUI", "custom_nodes")
 os.makedirs(custom_nodes_dir, exist_ok=True)
 
-# Clone UltimateSDUpscaler
-usd_dir = os.path.join(custom_nodes_dir, "ComfyUI_UltimateSDUpscaler")
+# Clone UltimateSDUpscale
+usd_dir = os.path.join(custom_nodes_dir, "ComfyUI_UltimateSDUpscale")
 if not os.path.exists(usd_dir):
-    subprocess.run(["git", "clone", "https://github.com/ssitu/ComfyUI_UltimateSDUpscaler", usd_dir])
+    subprocess.run(["git", "clone", "https://github.com/ssitu/ComfyUI_UltimateSDUpscale", usd_dir])
 
 # Clone comfy_mtb
 mtb_dir = os.path.join(custom_nodes_dir, "comfy_mtb")
@@ -47,11 +48,11 @@ if not os.path.exists(kjn_dir):
 comfy_models_dir = os.path.join("ComfyUI", "models")
 os.makedirs(comfy_models_dir, exist_ok=True)
 
-# UNET (Flux FP8)
-unet_dir = os.path.join(comfy_models_dir, "unet")
-os.makedirs(unet_dir, exist_ok=True)
-if not os.path.exists(os.path.join(unet_dir, "flux1-dev-fp8.safetensors")):
-    hf_hub_download(repo_id="Kijai/flux-fp8", filename="flux1-dev-fp8.safetensors", local_dir=unet_dir)
+# Diffusion models (Flux FP8)
+diffusion_dir = os.path.join(comfy_models_dir, "diffusion_models")
+os.makedirs(diffusion_dir, exist_ok=True)
+if not os.path.exists(os.path.join(diffusion_dir, "flux1-dev-fp8.safetensors")):
+    hf_hub_download(repo_id="Kijai/flux-fp8", filename="flux1-dev-fp8.safetensors", local_dir=diffusion_dir)
 
 # CLIP models
 clip_dir = os.path.join(comfy_models_dir, "clip")
@@ -204,8 +205,8 @@ def enhance_image(
         vaeloader = NODE_CLASS_MAPPINGS["VAELoader"]()
         vae = get_value_at_index(vaeloader.load_vae(vae_name="ae.safetensors"), 0)
 
-        unetloader = NODE_CLASS_MAPPINGS["UNETLoader"]()
-        model = get_value_at_index(unetloader.load_unet(unet_name="flux1-dev-fp8.safetensors", weight_dtype="fp8_e4m3fn"), 0)
+        unetloader = NODE_CLASS_MAPPINGS["LoadDiffusionModel"]()
+        model = get_value_at_index(unetloader.load_diffusion_model(unet_name="flux1-dev-fp8.safetensors", weight_dtype="fp8_e4m3fn"), 0)
 
         ultimatesdupscale = NODE_CLASS_MAPPINGS["UltimateSDUpscale"]()
         upscale_out = ultimatesdupscale.upscale(