Update app.py

app.py

@@ -1,103 +1,155 @@
+import torch
+from diffusers import FluxPipeline
 import gradio as gr
 import threading
 import os
-import torch
 
 os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
 torch.set_num_threads(os.cpu_count())
 
-model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
-model2 = gr.load("models/Purz/face-projection")
+# Initialize Flux pipeline
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16
+)
+pipe.enable_model_cpu_offload()
 
 stop_event = threading.Event()
 
-def generate_images(text, selected_model, steps, cfg_scale, seed, width, height):
+def generate_images(
+    prompt,
+    height,
+    width,
+    guidance_scale,
+    num_inference_steps,
+    max_sequence_length,
+    seed,
+    randomize_seed
+):
     stop_event.clear()
-
-    if selected_model == "Model 1 (Turbo Realism)":
-        model = model1
-    elif selected_model == "Model 2 (Face Projection)":
-        model = model2
-    else:
-        return ["Invalid model selection."] * 3
-
-    # Convert seed to integer (handle empty/None)
-    try:
-        seed = int(seed) if seed not in [None, ""] else -1
-    except:
-        seed = -1
-
     results = []
+    
     for i in range(3):
         if stop_event.is_set():
-            return ["Image generation stopped by user."] * 3
-
-        modified_text = f"{text} variation {i+1}"
-        result = model(
-            modified_text,
-            #num_inference_steps=int(steps),
-            #guidance_scale=float(cfg_scale),
+            return [None] * 3
+        
+        # Handle seed randomization
+        if randomize_seed:
+            current_seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        else:
+            current_seed = seed + i
+        
+        generator = torch.Generator(device="cpu").manual_seed(current_seed)
+        
+        # Generate image with current parameters
+        image = pipe(
+            prompt=prompt,
             height=int(height),
             width=int(width),
-            seed=seed if seed != -1 else None
-        )
-        results.append(result)
-
+            guidance_scale=guidance_scale,
+            num_inference_steps=int(num_inference_steps),
+            max_sequence_length=int(max_sequence_length),
+            generator=generator
+        ).images[0]
+        
+        results.append(image)
+    
     return results
 
 def stop_generation():
-    """Stops the ongoing image generation by setting the stop_event flag."""
     stop_event.set()
-    return ["Generation stopped."] * 3
+    return [None] * 3
 
 with gr.Blocks() as interface:
-    gr.Markdown(
-        "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
-    )
+    gr.Markdown("""
+    ### FLUX Image Generation
+    Adjust parameters below to control the image generation process
+    """)
     
     with gr.Row():
-        text_input = gr.Textbox(label="Prompt", placeholder="Type your imagination here...")
-        model_selector = gr.Radio(
-            ["Model 1 (Turbo Realism)", "Model 2 (Face Projection)"],
-            label="Model Selection",
-            value="Model 1 (Turbo Realism)"
+        text_input = gr.Textbox(
+            label="Prompt",
+            placeholder="Describe what you want to generate...",
+            scale=3
         )
     
-    with gr.Accordion("Advanced Parameters", open=False):
+    with gr.Accordion("Generation Parameters", open=False):
         with gr.Row():
-            steps = gr.Slider(1, 150, value=25, label="Inference Steps", info="(20-50 for quality/speed balance)")
-            cfg_scale = gr.Slider(1.0, 20.0, value=7.5, label="CFG Scale", info="(7-12 for good balance)")
-            seed = gr.Number(label="Seed", value=-1, precision=0, info="-1 for random")
+            height = gr.Number(
+                label="Height",
+                value=1024,
+                minimum=512,
+                maximum=4096,
+                step=64,
+                precision=0
+            )
+            width = gr.Number(
+                label="Width",
+                value=1024,
+                minimum=512,
+                maximum=4096,
+                step=64,
+                precision=0
+            )
+        
+        guidance_scale = gr.Slider(
+            label="Guidance Scale",
+            minimum=0.0,
+            maximum=20.0,
+            value=7.0,
+            step=0.5
+        )
+        
+        num_inference_steps = gr.Slider(
+            label="Inference Steps",
+            minimum=10,
+            maximum=150,
+            value=50,
+            step=1
+        )
+        
+        max_sequence_length = gr.Dropdown(
+            label="Max Sequence Length",
+            choices=[512, 768, 1024],
+            value=512
+        )
         
         with gr.Row():
-            width = gr.Dropdown(
-                choices=["512", "640", "768", "896", "1024"],
-                value="512",
-                label="Width",
-                allow_custom_value=True
+            seed = gr.Number(
+                label="Seed",
+                value=42,
+                precision=0
             )
-            height = gr.Dropdown(
-                choices=["512", "640", "768", "896", "1024"],
-                value="512",
-                label="Height",
-                allow_custom_value=True
+            randomize_seed = gr.Checkbox(
+                label="Randomize Seed",
+                value=True
             )
-    
+
     with gr.Row():
-        generate_button = gr.Button("Generate 3 Images 🎨", variant="primary")
-        stop_button = gr.Button("Stop Generation", variant="stop")
-    
+        generate_btn = gr.Button("Generate", variant="primary")
+        stop_btn = gr.Button("Stop Generation")
+
     with gr.Row():
-        output1 = gr.Image(label="Variant 1", type="pil", show_label=True)
-        output2 = gr.Image(label="Variant 2", type="pil", show_label=True)
-        output3 = gr.Image(label="Variant 3", type="pil", show_label=True)
+        output1 = gr.Image(label="Output 1", type="pil")
+        output2 = gr.Image(label="Output 2", type="pil")
+        output3 = gr.Image(label="Output 3", type="pil")
 
-    generate_button.click(
+    generate_btn.click(
         generate_images,
-        inputs=[text_input, model_selector, steps, cfg_scale, seed, width, height],
+        inputs=[
+            text_input,
+            height,
+            width,
+            guidance_scale,
+            num_inference_steps,
+            max_sequence_length,
+            seed,
+            randomize_seed
+        ],
         outputs=[output1, output2, output3]
     )
-    stop_button.click(
+    
+    stop_btn.click(
         stop_generation,
         inputs=[],
         outputs=[output1, output2, output3]