use huggingface-cli to download

app.py

@@ -23,6 +23,51 @@ login(token=os.getenv('HF_TOKEN'))
 device = torch.device('cuda')
 
 
+
+# Load models outside the function to avoid reloading every time
+def load_models():
+    # Load sd35 model
+    nf4_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+    model_nf4 = SD3Transformer2DModel.from_pretrained(
+        "stabilityai/stable-diffusion-3.5-large",
+        subfolder="transformer",
+        quantization_config=nf4_config,
+        torch_dtype=torch.bfloat16
+    )
+    pipe_sd35 = StableDiffusion3Pipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3.5-large", 
+        transformer=model_nf4,
+        torch_dtype=torch.bfloat16,
+    )
+    pipe_sd35.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe_sd35.scheduler.config)
+    inverse_scheduler_sd35 = FlowMatchEulerInverseScheduler.from_pretrained(
+        "stabilityai/stable-diffusion-3.5-large",
+        subfolder='scheduler'
+    )
+    pipe_sd35.inv_scheduler = inverse_scheduler_sd35
+
+    # Load sdxl model
+    pipe_sdxl = StableDiffusionXLPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0", 
+        torch_dtype=torch.float16, 
+        variant="fp16",
+        use_safetensors=True
+    ).to("cuda")
+    pipe_sdxl.scheduler = DDIMScheduler.from_config(pipe_sdxl.scheduler.config)
+    inverse_scheduler_sdxl = DDIMInverseScheduler.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        subfolder='scheduler'
+    )
+    pipe_sdxl.inv_scheduler = inverse_scheduler_sdxl
+
+    return pipe_sd35, pipe_sdxl
+
+pipe_sd35, pipe_sdxl = load_models()
+
 @spaces.GPU
 def generate_image(
     model_name,
@@ -41,71 +86,35 @@ def generate_image(
         torch.cuda.empty_cache()
         dtype = torch.float16
         set_seed(seed)
+
+        # Select the appropriate pipeline
         if model_name == 'sd35':
-            nf4_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.bfloat16
-            )
-            model_nf4 = SD3Transformer2DModel.from_pretrained(
-                "stabilityai/stable-diffusion-3.5-large",
-                subfolder="transformer",
-                quantization_config=nf4_config,
-                torch_dtype=torch.bfloat16
-            )
-
-            pipe = StableDiffusion3Pipeline.from_pretrained(
-                "stabilityai/stable-diffusion-3.5-large", 
-                transformer=model_nf4,
-                torch_dtype=torch.bfloat16,
-            )
-
-            pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
-            inverse_scheduler = FlowMatchEulerInverseScheduler.from_pretrained("stabilityai/stable-diffusion-3.5-large",
-                                                                  subfolder='scheduler')
-            pipe.inv_scheduler = inverse_scheduler
-
-        elif model_name == "sdxl":
-            pipe = StableDiffusionXLPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-xl-base-1.0", 
-                torch_dtype=torch.float16, 
-                variant="fp16",
-                use_safetensors=True
-            ).to("cuda")
-
-            pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-            inverse_scheduler = DDIMInverseScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
-                                                                  subfolder='scheduler')
-            pipe.inv_scheduler = inverse_scheduler
+            pipe = pipe_sd35
+        elif model_name == 'sdxl':
+            pipe = pipe_sdxl
+        else:
+            raise ValueError("Invalid model name")
 
         pipe.to(device)
         pipe.enable_model_cpu_offload()
-
+      #   os.makedirs('./weights', exist_ok=True)
+        os.system('huggingface-cli download sst12345/CoRe2 weights/sd35_noise_model.pth weights/sdxl_noise_model.pth --local-dir ./weights')
         # TODO: load noise model
         if method == 'core' or method == 'z-core':
             from diffusion_pipeline.refine_model import PromptSD35Net, PromptSDXLNet
             from diffusion_pipeline.lora import replace_linear_with_lora, lora_true
 
             if model_name == 'sd35':
-                refine_model = PromptSD35Net()
+                refine_model = PromptSD35Net().to(device)
                 replace_linear_with_lora(refine_model, rank=64, alpha=1.0, number_of_lora=28)
                 lora_true(refine_model, lora_idx=0)
-
-                os.makedirs('./weights', exist_ok=True)
-                if not os.path.exists('./weights/sd35_noise_model.pth'):
-                    os.system('wget https://huggingface.co/sst12345/CoRe2/resolve/main/weights/sd35_noise_model.pth')
-                    os.system('mv sd35_noise_model.pth ./weights/')
-                checkpoint = torch.load('./weights/sd35_noise_model.pth', map_location='cpu')
+                checkpoint = torch.load('./weights/weights/sd35_noise_model.pth', map_location='cpu')
                 refine_model.load_state_dict(checkpoint)
             elif model_name == 'sdxl':
                 refine_model = PromptSDXLNet()
                 replace_linear_with_lora(refine_model, rank=48, alpha=1.0, number_of_lora=50)
                 lora_true(refine_model, lora_idx=0)
-                os.makedirs('./weights', exist_ok=True)
-                if not os.path.exists('./weights/sdxl_noise_model.pth'):
-                    os.system('wget https://huggingface.co/sst12345/CoRe2/resolve/main/weights/sdxl_noise_model.pth')
-                    os.system('mv sdxl_noise_model.pth ./weights/')
-                checkpoint = torch.load('./weights/sdxl_noise_model.pth', map_location='cpu')
+                checkpoint = torch.load('./weights/weights/sdxl_noise_model.pth', map_location='cpu')
                 refine_model.load_state_dict(checkpoint)
 
             print("Load Lora Success")
@@ -235,7 +244,7 @@ if __name__ == '__main__':
             gr.Slider(minimum=1024, maximum=2048, value=1024, label="Size")  # 设置默认大小为 1024
       ],
       outputs=gr.Image(type="filepath"),  # 修改了type参数
-      title="Image Generation with CoRe^2"
+      title="Image Generation with CoRe^2",
       )
-      iface.launch()
+      iface.launch(share=True)
     

requirements.txt

@@ -1,12 +1,13 @@
-diffusers
-transformers
+diffusers==0.31.0
+transformers==4.46.1
 einops 
 wandb
-accelerate
+accelerate==1.0.1
 pandas
 imageio
 gradio
 sentencepiece
+huggingface_hub
 bitsandbytes
 imageio-ffmpeg
 omegaconf