Switch to FLUX.1-schnell model

app.py

@@ -1,63 +1,55 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionPipeline
-from diffusers.utils import load_image
+from diffusers import FluxPipeline, FluxImg2ImgPipeline
 
-print("Loading Models...")
-# 1. Text-to-Image Model (Standard Generation)
-pipe_t2i = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16,
-    safety_checker=None
-)
-
-# 2. Instruct-Pix2Pix Model (Instruction-Based Editing)
-pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
-    "timbrooks/instruct-pix2pix",
-    torch_dtype=torch.float16,
-    safety_checker=None
+print("Loading FLUX.1 Models (CPU)...")
+# Load the base FLUX model
+# We load the pipelines on CPU so ZeroGPU doesn't crash during startup.
+pipe_t2i = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
 )
 
+pipe_edit = FluxImg2ImgPipeline.from_pipe(pipe_t2i)
 print("Models loaded successfully!")
 
 @spaces.GPU
-def generate_or_edit(prompt, input_image, image_guidance_scale, guidance_scale, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+def generate_or_edit(prompt, input_image, denoising_strength, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+    # Move pipelines to GPU inside the ZeroGPU decorated function
     pipe_t2i.to("cuda")
     pipe_edit.to("cuda")
+    
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
     generator = torch.Generator("cuda").manual_seed(int(seed))
     
     if not prompt:
-        raise gr.Error("Please enter a prompt or instruction!")
+        raise gr.Error("Please enter a prompt!")
     
     if input_image is not None:
-        # Edit mode
+        # Edit mode (Img2Img)
         input_image = input_image.convert("RGB")
-        # Resize image for SD1.5 (InstructPix2Pix)
-        input_image = input_image.resize((512, 512))
-        
         image = pipe_edit(
-            prompt,
+            prompt=prompt,
             image=input_image,
+            strength=denoising_strength,
             num_inference_steps=int(num_inference_steps),
-            image_guidance_scale=image_guidance_scale,
-            guidance_scale=guidance_scale,
+            guidance_scale=0.0, # FLUX.1-schnell uses 0 guidance scale
             generator=generator,
         ).images[0]
     else:
-        # Generate mode
+        # Generate mode (T2I)
         image = pipe_t2i(
-            prompt,
+            prompt=prompt,
             num_inference_steps=int(num_inference_steps),
-            guidance_scale=guidance_scale,
+            guidance_scale=0.0, # FLUX.1-schnell uses 0 guidance scale
             generator=generator,
         ).images[0]
         
     return image, seed
 
-# Build the Gradio interface
+# UI
 custom_theme = gr.themes.Soft(
     primary_hue="blue",
     secondary_hue="indigo",
@@ -67,34 +59,40 @@ custom_theme = gr.themes.Soft(
 with gr.Blocks(theme=custom_theme, fill_height=True) as demo:
     gr.Markdown(
         """
-        # 🎨 AI Image Studio (Generation & Instruction Editing)
-        Generate images from scratch, or upload an image and use a prompt like *"Make him wear sunglasses"* to edit it seamlessly.
+        # ⚡ FLUX.1 Image Studio (Grok Quality)
+        Generate state-of-the-art images from scratch, or edit existing ones using the FLUX.1-schnell model.
         """
     )
     
     with gr.Row():
         with gr.Column(scale=1):
             prompt = gr.Textbox(
-                label="✨ Instruction / Prompt", 
+                label="✨ Prompt", 
                 lines=3, 
-                placeholder="e.g. 'A futuristic city at night' (for generation) OR 'Turn the daytime into nighttime' (for editing)",
+                placeholder="e.g. A cyberpunk cat...", 
                 autofocus=True
             )
             input_image = gr.Image(
-                label="🖼️ Input Image (Optional - Leave blank to generate from scratch)", 
+                label="🖼️ Input Image (Optional - For editing)", 
                 type="pil"
             )
             
             with gr.Accordion("⚙️ Advanced Settings", open=False):
-                num_inference_steps = gr.Slider(minimum=10, maximum=50, value=30, step=1, label="Inference Steps")
-                guidance_scale = gr.Slider(minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Text Guidance Scale")
-                image_guidance_scale = gr.Slider(
-                    minimum=1.0, 
-                    maximum=3.0, 
-                    value=1.5, 
-                    step=0.1, 
-                    label="Image Guidance Scale (Editing Only)", 
-                    info="Lower values = more changes. Higher values = stays closer to original image."
+                denoising_strength = gr.Slider(
+                    minimum=0.0, 
+                    maximum=1.0, 
+                    value=0.5, 
+                    step=0.05, 
+                    label="Denoising Strength (Editing Only)", 
+                    info="Lower = keeps more of original image. Higher = completely changes image to match prompt."
+                )
+                num_inference_steps = gr.Slider(
+                    minimum=1, 
+                    maximum=12, 
+                    value=4, 
+                    step=1, 
+                    label="Inference Steps", 
+                    info="FLUX.1-schnell is optimized for 4 steps."
                 )
                 
                 with gr.Row():
@@ -113,15 +111,14 @@ with gr.Blocks(theme=custom_theme, fill_height=True) as demo:
             output_image = gr.Image(label="Result", type="pil", interactive=False)
             used_seed = gr.Number(label="Seed Used", interactive=False)
 
-    # Connections
     generate_btn.click(
         fn=generate_or_edit,
-        inputs=[prompt, input_image, image_guidance_scale, guidance_scale, num_inference_steps, seed, randomize_seed],
+        inputs=[prompt, input_image, denoising_strength, num_inference_steps, seed, randomize_seed],
         outputs=[output_image, used_seed]
     )
     prompt.submit(
         fn=generate_or_edit,
-        inputs=[prompt, input_image, image_guidance_scale, guidance_scale, num_inference_steps, seed, randomize_seed],
+        inputs=[prompt, input_image, denoising_strength, num_inference_steps, seed, randomize_seed],
         outputs=[output_image, used_seed]
     )