Update app.py

app.py

@@ -2,12 +2,19 @@ import torch
 from diffusers import FluxPipeline
 from huggingface_hub import login
 
-# Replace with your Hugging Face token (get it from: https://huggingface.co/settings/tokens)
-login(token="hf_yourtokenhere")
+# Authenticate (required for gated models)
+login(token="hf_yourtokenhere")  # Replace with your token
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+# Load the model (use bfloat16 for faster inference + less VRAM)
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16
+)
 
+# Enable CPU offloading if you have limited GPU memory
+pipe.enable_model_cpu_offload()
+
+# Generate an image
 prompt = "A cat holding a sign that says hello world"
 image = pipe(
     prompt,
@@ -16,59 +23,7 @@ image = pipe(
     max_sequence_length=256,
     generator=torch.Generator("cpu").manual_seed(0)
 ).images[0]
-image.save("flux-schnell.png")
-import torch
-from diffusers import FluxPipeline
-import gradio as gr
-
-def generate_image(prompt, guidance_scale, num_inference_steps, seed):
-    # Load the pipeline (you might want to move this outside the function for better performance)
-    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-    pipe.enable_model_cpu_offload()
-    
-    # Generate the image
-    generator = torch.Generator("cpu").manual_seed(seed)
-    image = pipe(
-        prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        max_sequence_length=256,
-        generator=generator
-    ).images[0]
-    
-    return image
-
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# FLUX.1-schnell Image Generator")
-    
-    with gr.Row():
-        with gr.Column():
-            prompt = gr.Textbox(label="Prompt", value="A cat holding a sign that says hello world")
-            guidance_scale = gr.Slider(minimum=0.0, maximum=20.0, value=0.0, label="Guidance Scale")
-            num_steps = gr.Slider(minimum=1, maximum=20, step=1, value=4, label="Inference Steps")
-            seed = gr.Number(value=0, label="Seed", precision=0)
-            submit_btn = gr.Button("Generate Image")
-        
-        with gr.Column():
-            output_image = gr.Image(label="Generated Image")
-    
-    submit_btn.click(
-        fn=generate_image,
-        inputs=[prompt, guidance_scale, num_steps, seed],
-        outputs=output_image
-    )
-    
-    gr.Examples(
-        examples=[
-            ["A futuristic cityscape at sunset, digital art", 0.0, 4, 42],
-            ["An astronaut riding a horse on Mars, photorealistic", 0.0, 6, 123],
-            ["A steampunk library with floating books", 0.0, 4, 7]
-        ],
-        inputs=[prompt, guidance_scale, num_steps, seed],
-        outputs=output_image,
-        fn=generate_image,
-        cache_examples=True
-    )
 
-demo.launch()
+# Save the output
+image.save("flux-schnell-output.png")
+print("Image saved successfully!")