Update app.py

app.py

@@ -8,7 +8,7 @@ import numpy as np
 
 # Set data type
 dtype = torch.bfloat16
-device = "cpu"  # Initial device for model loading; inference will use GPU
+device = "cpu"  # Use CPU for model loading to avoid CUDA initialization
 
 # Load tokenizer and text encoder for Llama
 try:
@@ -25,12 +25,12 @@ except Exception as e:
 # Load the HiDreamImagePipeline
 try:
     pipe = HiDreamImagePipeline.from_pretrained(
-        "HiDream-ai/HiDream-I1-Full",
+        "HiDream-ai/HiDream-I1-Dev",
         tokenizer_4=tokenizer_4,
         text_encoder_4=text_encoder_4,
         torch_dtype=dtype,
     ).to(device)
-    pipe.enable_model_cpu_offload()  # Offload to CPU when not in use, critical for Spaces
+    pipe.enable_model_cpu_offload()  # Offload to CPU, automatically manages GPU placement
 except Exception as e:
     raise Exception(f"Failed to load HiDreamImagePipeline: {e}. Ensure you have access to 'HiDream-ai/HiDream-I1-Full'.")
 
@@ -40,16 +40,13 @@ MAX_IMAGE_SIZE = 2048
 
 # Inference function with GPU access
 @spaces.GPU()
-def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=5.0, progress=gr.Progress(track_tqdm=True)):
-    # Ensure the model is on GPU for inference
-    pipe.to("cuda")
-    
+def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=28, guidance_scale=3.5, progress=gr.Progress(track_tqdm=True)):
     try:
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         generator = torch.Generator("cuda").manual_seed(seed)
         
-        # Generate the image
+        # Generate the image; offloading handles device placement
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -61,13 +58,9 @@ def infer(prompt, negative_prompt="", seed=42, randomize_seed=False, width=1024,
             output_type="pil",
         ).images[0]
         
-        # Clear GPU memory
-        torch.cuda.empty_cache()
-        
         return image, seed
     finally:
-        # Move model back to CPU to free GPU memory
-        pipe.to("cpu")
+        # Clear GPU memory
         torch.cuda.empty_cache()
 
 # Define examples
@@ -89,7 +82,7 @@ css = """
    color: white !important;
 }
 .generate-btn:hover {
-   transform: translateY(-2px);
+   transform: translateY2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
 }
 """
@@ -107,6 +100,12 @@ with gr.Blocks(css=css) as app:
                         lines=3,
                         elem_id="prompt-text-input"
                     )
+                with gr.Row():
+                    negative_prompt = gr.Textbox(
+                        label="Negative Prompt",
+                        placeholder="Enter what to avoid (optional)",
+                        lines=2
+                    )
                 with gr.Row():
                     with gr.Accordion("Advanced Settings", open=False):
                         with gr.Row():
@@ -127,14 +126,14 @@ with gr.Blocks(css=css) as app:
                         with gr.Row():
                             steps = gr.Slider(
                                 label="Inference Steps",
-                                value=50,
+                                value=28,
                                 minimum=1,
                                 maximum=100,
                                 step=1
                             )
                             cfg = gr.Slider(
                                 label="Guidance Scale",
-                                value=5.0,
+                                value=3.5,
                                 minimum=1,
                                 maximum=20,
                                 step=0.5
@@ -151,12 +150,6 @@ with gr.Blocks(css=css) as app:
                                 label="Randomize Seed",
                                 value=True
                             )
-                        with gr.Row():
-                            negative_prompt = gr.Textbox(
-                                label="Negative Prompt",
-                                placeholder="Enter what to avoid (optional)",
-                                lines=2
-                            )
                 with gr.Row():
                     text_button = gr.Button(
                         "✨ Generate Image",
@@ -189,5 +182,4 @@ with gr.Blocks(css=css) as app:
         outputs=[image_output, seed_output]
     )
 
-# Launch the app
 app.launch(share=True)