ZeroGPU fix

app.py

@@ -7,11 +7,13 @@ from diffusers.pipelines.flux.pipeline_flux_controlnet import FluxControlNetPipe
 from diffusers.models.controlnet_flux import FluxControlNetModel
 import numpy as np
 from huggingface_hub import login, snapshot_download
+import spaces
+
 
 # Configuration
-base_model = 'black-forest-labs/FLUX.1-dev'
-controlnet_model = 'promeai/FLUX.1-controlnet-lineart-promeai'
-css = """
+BASE_MODEL = 'black-forest-labs/FLUX.1-dev'
+CONTROLNET_MODEL = 'promeai/FLUX.1-controlnet-lineart-promeai'
+CSS = """
 #col-container {
     margin: 0 auto;
     max-width: 640px;
@@ -19,28 +21,28 @@ css = """
 """
 
 # Setup
-auth_token = os.getenv("HF_AUTH_TOKEN")
-if not auth_token:
+AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
+if AUTH_TOKEN:
+    login(AUTH_TOKEN)
+else:
     raise ValueError("Hugging Face auth token not found. Please set HF_AUTH_TOKEN in the environment.")
 
-login(auth_token)
-
-model_dir = snapshot_download(
-    repo_id=base_model,
+MODEL_DIR = snapshot_download(
+    repo_id=BASE_MODEL,
     revision="main",
-    use_auth_token=auth_token
+    use_auth_token=AUTH_TOKEN
 )
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-print(f"Using device: {device} (torch_dtype={torch_dtype})")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TORCH_DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch_dtype)
-pipe = FluxControlNetPipeline.from_pretrained(model_dir, controlnet=controlnet, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+CONTROLNET = FluxControlNetModel.from_pretrained(CONTROLNET_MODEL, torch_dtype=TORCH_DTYPE)
+PIPE = FluxControlNetPipeline.from_pretrained(MODEL_DIR, controlnet=CONTROLNET, torch_dtype=TORCH_DTYPE)
+PIPE = PIPE.to(DEVICE)
 
 MAX_SEED = np.iinfo(np.int32).max
 
+@spaces.GPU
 def infer(
     prompt,
     control_image_path,
@@ -50,6 +52,9 @@ def infer(
     seed,
     randomize_seed,
 ):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    print(f"Inference: using device: {device} (torch_dtype={torch_dtype})")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -57,7 +62,7 @@ def infer(
     control_image = load_image(control_image_path) if control_image_path else None
 
     # Generate image
-    result = pipe(
+    result = PIPE(
         prompt=prompt,
         control_image=control_image,
         controlnet_conditioning_scale=controlnet_conditioning_scale,
@@ -68,7 +73,7 @@ def infer(
 
     return result, seed
 
-with gr.Blocks(css=css) as demo:
+with gr.Blocks(css=CSS) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("Flux.1[dev] LineArt")
         gr.Markdown("### Zero-shot Partial Style Transfer for Line Art Images, Powered by FLUX.1")
@@ -118,7 +123,8 @@ with gr.Blocks(css=css) as demo:
         gr.Examples(
             examples=[
                 "Shiba Inu wearing dinosaur costume riding skateboard",
-                "Victorian style mansion interior with candlelight"
+                "Victorian style mansion interior with candlelight",
+                "Loading screen for Grand Theft Otter: Clam Andreas"
             ],
             inputs=[prompt]
         )