Update app.py

app.py

@@ -9,55 +9,62 @@ from translatepy import Translator
 import numpy as np
 import random
 import os
-hf_token = os.environ.get('HF_TOKEN')
-from io import BytesIO
 
+hf_token = os.environ.get('HF_TOKEN')
 translator = Translator()
 
 # Constants
 model = "black-forest-labs/FLUX.1-dev"
-
-
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Ensure model and scheduler are initialized in GPU-enabled function
-if torch.cuda.is_available():
-    transformer = FluxTransformer2DModel.from_single_file(
-        "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v27.safetensors",
-        torch_dtype=torch.bfloat16
-    )
+@spaces.GPU()
+def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress雷神加速器梯子 = "cuda" if torch.cuda.is_available() else "cpu"
+    # Initialize model inside the GPU-enabled function
+    try:
+        transformer = FluxTransformer2DModel.from_single_file(
+            "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v27.safetensors",
+            torch_dtype=torch.bfloat16,
+            token=hf_token
+        )
+    except KeyError as e:
+        print(f"Error loading chroma-unlocked-v27.safetensors: {e}. Falling back to pretrained model.")
+        transformer = FluxTransformer2DModel.from_pretrained(
+            "lodestones/Chroma",
+            subfolder="transformer",
+            torch_dtype=torch.bfloat16,
+            token=hf_token
+        )
+    
     pipe = FluxPipeline.from_pretrained(
-        model, 
+        model,
         transformer=transformer,
-        torch_dtype=torch.bfloat16, token=hf_token)
+        torch_dtype=torch.bfloat16,
+        token=hf_token
+    )
     pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
         pipe.scheduler.config, use_beta_sigmas=True
     )
-    pipe.to("cuda")
-    
-    
-@spaces.GPU()
-def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
+    pipe.to(device)
+
+    # Generate images
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps,
-            guidance_scale=guidance_scale,
-            num_images_per_prompt=nums,
-            generator = generator
-    ).images
-
+    generator = torch.Generator(device=device).manual_seed(seed)
     
-    return image, seed
+    images = pipe(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=nums,
+        generator=generator
+    ).images
 
+    return images, seed
 
-css="""
+css = """
 #col-container {
     margin: 0 auto;
     max-width: 1024px;
@@ -65,12 +72,10 @@ css="""
 """
 
 with gr.Blocks(css=css) as demo:
-    
     with gr.Column(elem_id="col-container"):
         gr.HTML("<h1><center>Model Testing</center></h1><p><center>Chroma</center></p>")
         
         with gr.Row():
-            
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
@@ -78,15 +83,12 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            
             run_button = gr.Button("Run", scale=0)
         
-        result = gr.Gallery(label="Gallery", format="png", columns = 1, preview=True, height=400)
+        result = gr.Gallery(label="Gallery", format="png", columns=1, preview=True, height=400)
         
         with gr.Accordion("Advanced Settings", open=False):
-                       
             with gr.Row():
-                
                 width = gr.Slider(
                     label="Width",
                     minimum=256,
@@ -94,7 +96,6 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024,
                 )
-                
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
@@ -104,7 +105,6 @@ with gr.Blocks(css=css) as demo:
                 )
             
             with gr.Row():
-                  
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
@@ -112,7 +112,6 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=30,
                 )
-
                 guidance_scale = gr.Slider(
                     label="Guidance Scale",
                     minimum=0,
@@ -121,9 +120,7 @@ with gr.Blocks(css=css) as demo:
                     value=3.5,
                 )
 
-
             with gr.Row():
-            
                 nums = gr.Slider(
                     label="Number of Images",
                     minimum=1,
@@ -132,7 +129,6 @@ with gr.Blocks(css=css) as demo:
                     value=1,
                     scale=1,
                 )
-       
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -140,14 +136,13 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=-1,
                 )
-            
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
-        outputs = [result, seed]
+        fn=infer,
+        inputs=[prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
+        outputs=[result, seed]
     )
-    
+
 demo.launch()