Update app.py

app.py

@@ -2,42 +2,58 @@ import os
 import random
 import gradio as gr
 import torch
-from diffusers import StableDiffusionPipeline, LCMScheduler
-from huggingface_hub import hf_hub_download
+from diffusers import DiffusionPipeline, UNet2DConditionModel
+from transformers import CLIPTextModel, CLIPTokenizer
 
-# Configuration
-MODEL_ID = "segmind/SSD-1B"
-LORA_ID = "latent-consistency/lcm-lora-ssd-1b"
+# Configuration - Using Flux Model
+MODEL_ID = "CompVis/Flux-Pro"
+LORA_ID = "flux/lora-weights"
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
-# Initialize the pipeline with optimizations
 def get_pipeline():
-    # Load base model
-    pipe = StableDiffusionPipeline.from_pretrained(
+    # Load Flux components
+    unet = UNet2DConditionModel.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.float32,
+        subfolder="unet",
         cache_dir=MODEL_CACHE,
-        safety_checker=None,
-        use_safetensors=True
+        torch_dtype=torch.float32
+    )
+    
+    text_encoder = CLIPTextModel.from_pretrained(
+        MODEL_ID,
+        subfolder="text_encoder",
+        cache_dir=MODEL_CACHE
     )
     
-    # Apply LCM-LoRA for 4-8x faster generation
-    lcm_lora_path = hf_hub_download(
-        LORA_ID, 
-        "lora_layers.safetensors",
+    tokenizer = CLIPTokenizer.from_pretrained(
+        MODEL_ID,
+        subfolder="tokenizer",
         cache_dir=MODEL_CACHE
     )
-    pipe.load_lora_weights(lcm_lora_path)
-    pipe.fuse_lora()
     
-    # Configure scheduler
-    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    # Create pipeline
+    pipe = DiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        unet=unet,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        cache_dir=MODEL_CACHE,
+        torch_dtype=torch.float32,
+        safety_checker=None
+    )
+    
+    # Load LoRA weights
+    lora_path = hf_hub_download(
+        LORA_ID,
+        "flux_lora.safetensors",
+        cache_dir=MODEL_CACHE
+    )
+    pipe.unet.load_attn_procs(lora_path)
     
     # CPU optimizations
     pipe = pipe.to("cpu")
     pipe.enable_attention_slicing()
-    pipe.enable_model_cpu_offload()
     
     return pipe
 
@@ -50,15 +66,13 @@ def generate_image(
     width: int = 768,
     height: int = 768,
     seed: int = -1,
-    guidance_scale: float = 2.0,  # Lower for LCM
-    num_inference_steps: int = 8   # Only 4-8 steps needed
+    guidance_scale: float = 7.5,
+    num_inference_steps: int = 25
 ):
-    # Handle seed
     if seed == -1:
         seed = random.randint(0, 2147483647)
     generator = torch.Generator(device="cpu").manual_seed(seed)
     
-    # Generate image
     with torch.no_grad():
         image = pipeline(
             prompt=prompt,
@@ -67,118 +81,36 @@ def generate_image(
             height=height,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
-            generator=generator,
-            output_type="pil"
+            generator=generator
         ).images[0]
     
     return image, seed
 
-# Professional enhancement presets
-ENHANCEMENTS = {
-    "Ultra Realistic": "masterpiece, best quality, ultra detailed, 8k",
-    "Cinematic": "cinematic lighting, film grain, depth of field",
-    "Digital Art": "concept art, digital painting, sharp focus",
-    "Photorealistic": "photorealistic, DSLR, f/1.8, natural lighting"
-}
-
-def apply_enhancement(prompt, enhancement):
-    return f"{ENHANCEMENTS[enhancement]}, {prompt}"
-
-# Create professional interface
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal")) as demo:
-    gr.Markdown("""
-    # 🚀 FLUX Pro Image Generator
-    **Professional Quality Images · Lightning Fast on CPU**  
-    *Powered by SSD-1B with LCM-LoRA technology*
-    """)
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🌀 FLUX-Pro Image Generator")
     
     with gr.Row():
-        with gr.Column(scale=3):
-            with gr.Group():
-                prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="Describe the image you want to create...",
-                    lines=3,
-                    elem_id="prompt-box"
-                )
-                enhancement = gr.Dropdown(
-                    label="Quality Enhancement",
-                    choices=list(ENHANCEMENTS.keys()),
-                    value="Ultra Realistic"
-                )
-                negative_prompt = gr.Textbox(
-                    label="Negative Prompt",
-                    placeholder="What to exclude from the image...",
-                    value="low quality, blurry, cartoon, drawing, sketch"
-                )
-            
-            generate_btn = gr.Button("Generate Image", variant="primary", size="lg")
+        with gr.Column():
+            prompt = gr.Textbox(label="Prompt", lines=3)
+            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, low quality")
+            generate_btn = gr.Button("Generate", variant="primary")
             
-            with gr.Accordion("Advanced Settings", open=False):
-                with gr.Row():
-                    width = gr.Slider(512, 1024, value=768, step=64, label="Width")
-                    height = gr.Slider(512, 1024, value=768, step=64, label="Height")
-                with gr.Row():
-                    guidance = gr.Slider(1.0, 5.0, value=2.0, step=0.1, label="Guidance Scale")
-                    steps = gr.Slider(4, 12, value=8, step=1, label="Generation Steps")
+            with gr.Accordion("Advanced", open=False):
+                width = gr.Slider(512, 1024, value=768, step=64, label="Width")
+                height = gr.Slider(512, 1024, value=768, step=64, label="Height")
+                guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="Guidance")
+                steps = gr.Slider(15, 50, value=25, step=1, label="Steps")
                 seed = gr.Number(label="Seed", value=-1)
         
-        with gr.Column(scale=2):
-            output_image = gr.Image(
-                label="Generated Image", 
-                type="pil",
-                elem_id="output-image",
-                height=500
-            )
-            with gr.Group():
-                used_seed = gr.Textbox(label="Used Seed", interactive=False)
-                gr.Markdown("**Tip:** For best results, use detailed descriptive prompts")
-    
-    # Enhancement application
-    enhancement.change(
-        fn=apply_enhancement,
-        inputs=[prompt, enhancement],
-        outputs=prompt
-    )
-    
-    # Generation handler
+        with gr.Column():
+            output_image = gr.Image(label="Result", type="pil")
+            used_seed = gr.Textbox(label="Used Seed")
+
     generate_btn.click(
-        fn=generate_image,
+        generate_image,
         inputs=[prompt, negative_prompt, width, height, seed, guidance, steps],
         outputs=[output_image, used_seed]
     )
-    
-    # Professional examples
-    gr.Examples(
-        examples=[
-            [
-                "A futuristic cityscape at twilight, towering skyscrapers with neon lights reflecting on wet streets", 
-                "Ultra Realistic",
-                "low quality, cartoon, drawing, blurry",
-                768,
-                768
-            ],
-            [
-                "Portrait of a cyberpunk samurai warrior, intricate armor with glowing circuits, detailed face", 
-                "Cinematic",
-                "anime, cartoon, sketch, simple background",
-                768,
-                1024
-            ],
-            [
-                "Majestic fantasy castle floating in the clouds, golden sunlight, hyperdetailed", 
-                "Digital Art",
-                "photograph, realistic, photo",
-                1024,
-                768
-            ]
-        ],
-        inputs=[prompt, enhancement, negative_prompt, width, height],
-        label="Professional Examples"
-    )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(
-        server_name="0.0.0.0",
-        server_port=int(os.getenv("PORT", 7860))
+    demo.launch()