Create app.py

app.py

@@ -0,0 +1,174 @@
+import streamlit as st
+import torch
+import numpy as np
+from diffusers import AutoPipelineForText2Image
+from huggingface_hub import login
+from tqdm.auto import tqdm
+from PIL import Image
+import os
+
+# Set page config
+st.set_page_config(
+    page_title="FLUX.1 Image Generator",
+    page_icon="🎨",
+    layout="wide"
+)
+
+# Custom CSS to improve the app's appearance
+st.markdown("""
+    <style>
+    .stProgress > div > div > div > div {
+        background-color: #1f77b4;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
+class StreamlitProgressCallback:
+    def __init__(self, progress_bar):
+        self.progress_bar = progress_bar
+        self.current_step = 0
+        
+    def __call__(self, step: int, timestep: int, latents: torch.FloatTensor):
+        self.current_step += 1
+        self.progress_bar.progress(self.current_step / step)
+
+@st.cache_resource
+def setup_flux(hf_token):
+    """
+    Setup FLUX.1 with proper authentication and GPU optimization
+    """
+    if not torch.cuda.is_available():
+        st.warning("⚠️ No GPU detected. Processing will be slow on CPU.")
+        return None
+
+    # Login to Hugging Face
+    login(token=hf_token)
+    
+    try:
+        model_id = "black-forest-labs/FLUX.1-dev"
+        with st.spinner(f"Loading {model_id}..."):
+            pipe = AutoPipelineForText2Image.from_pretrained(
+                model_id,
+                token=hf_token
+            )
+            
+        # Move to GPU and optimize
+        pipe = pipe.to("cuda")
+        pipe.enable_attention_slicing()
+        pipe.enable_vae_slicing()
+        
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+            st.success("✅ Model loaded successfully with xformers optimization!")
+        except Exception:
+            st.info("ℹ️ xformers not available. Install with: pip install xformers")
+        
+        return pipe
+        
+    except Exception as e:
+        st.error(f"Error loading FLUX.1: {str(e)}")
+        return None
+
+def process_image(image):
+    """Process the generated image to ensure proper value ranges"""
+    img_array = np.array(image)
+    img_array = np.clip(img_array, 0, 1)
+    img_array = (img_array * 255).round().astype(np.uint8)
+    return Image.fromarray(img_array)
+
+def generate_image(pipe, prompt, num_inference_steps=30):
+    """Generate an image using the initialized pipeline"""
+    if pipe is None:
+        st.error("Error: Pipeline not properly initialized")
+        return None
+        
+    # Clear CUDA cache before generation
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    
+    # Set up the progress callback
+    pipe.callback = StreamlitProgressCallback(progress_bar)
+    
+    try:
+        with torch.autocast("cuda"):
+            status_text.text("🎨 Generating image...")
+            image = pipe(
+                prompt=prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=7.5,
+            ).images[0]
+        
+        # Process the image
+        image = process_image(image)
+        status_text.text("✨ Generation complete!")
+        progress_bar.progress(1.0)
+        return image
+        
+    except Exception as e:
+        st.error(f"Error during image generation: {str(e)}")
+        return None
+    finally:
+        # Clean up GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+def main():
+    st.title("🎨 FLUX.1 Image Generator")
+    st.markdown("Generate amazing images using the FLUX.1 model from Black Forest Labs!")
+    
+    # Sidebar for configuration
+    with st.sidebar:
+        st.header("Configuration")
+        hf_token = st.text_input("HuggingFace Token", type="password", 
+                                help="Enter your HuggingFace token to access the model")
+        
+        num_steps = st.slider("Number of inference steps", 
+                             min_value=10, max_value=50, value=30,
+                             help="More steps generally means better quality but slower generation")
+        
+        st.markdown("---")
+        st.markdown("### Tips for better prompts:")
+        st.markdown("""
+        - Be specific and descriptive
+        - Include details about style, lighting, and mood
+        - Mention artistic mediums or techniques
+        """)
+    
+    # Main content area
+    prompt = st.text_area("Enter your prompt", 
+                         "A serene landscape with mountains and a lake at sunset",
+                         help="Describe the image you want to generate")
+    
+    col1, col2 = st.columns([1, 4])
+    with col1:
+        generate_button = st.button("Generate Image", type="primary")
+    
+    # Initialize session state for storing generated images
+    if 'generated_images' not in st.session_state:
+        st.session_state.generated_images = []
+    
+    if generate_button and hf_token:
+        pipe = setup_flux(hf_token)
+        if pipe is not None:
+            image = generate_image(pipe, prompt, num_steps)
+            if image is not None:
+                # Save the image
+                st.session_state.generated_images.append({
+                    'image': image,
+                    'prompt': prompt
+                })
+    
+    # Display generated images
+    if st.session_state.generated_images:
+        st.markdown("---")
+        st.header("Generated Images")
+        for idx, item in enumerate(reversed(st.session_state.generated_images)):
+            st.markdown(f"**Prompt:** {item['prompt']}")
+            st.image(item['image'], use_column_width=True)
+            st.markdown("---")
+
+if __name__ == "__main__":
+    main()