Update raw.py

raw.py

@@ -1,38 +1,55 @@
 import torch
-from diffusers.utils import load_image
-from diffusers import FluxControlNetModel
-from diffusers.pipelines import FluxControlNetPipeline
 import spaces
+import os
+from diffusers.utils import load_image
+from diffusers import FluxControlNetModel, FluxControlNetPipeline
+import gradio as gr
+huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
 # Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
-  "jasperai/Flux.1-dev-Controlnet-Upscaler",
-  torch_dtype=torch.bfloat16
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    torch_dtype=torch.bfloat16
 )
 pipe = FluxControlNetPipeline.from_pretrained(
-  "black-forest-labs/FLUX.1-dev",
-  controlnet=controlnet,
-  torch_dtype=torch.bfloat16
+    "black-forest-labs/FLUX.1-dev",
+    controlnet=controlnet,
+    torch_dtype=torch.bfloat16,
+    token=huggingface_token
 )
 pipe.to("cuda")
 
-# Load a control image
-control_image = load_image(
-  "https://img.getimg.ai/generated/img-HvTBfR1W1RUhykk7gmBUYp.png"
-)
+@spaces.GPU
+def generate_image(prompt, control_image):
+    # Load control image
+    control_image = load_image(control_image.name)
+    w, h = control_image.size
+    # Upscale x4
+    control_image = control_image.resize((w * 2, h * 2))
+    image = pipe(
+        prompt=prompt,
+        control_image=control_image,
+        controlnet_conditioning_scale=0.6,
+        num_inference_steps=14,
+        guidance_scale=3.5,
+        height=control_image.size[1],
+        width=control_image.size[0]
+    ).images[0]
+    return image
 
-w, h = control_image.size
-
-# Upscale x4
-control_image = control_image.resize((w * 2, h * 2))
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.inputs.Image(type="pil", label="Control Image"),
+    ],
+    outputs=[
+        gr.outputs.Image(type="pil", label="Generated Image"),
+    ],
+    title="FLUX ControlNet Image Generation",
+    description="Generate images using the FluxControlNetPipeline. Upload a control image and enter a prompt to create an image.",
+)
 
-image = pipe(
-    prompt="", 
-    control_image=control_image,
-    controlnet_conditioning_scale=0.6,
-    num_inference_steps=14, 
-    guidance_scale=3.5,
-    height=control_image.size[1],
-    width=control_image.size[0]
-).images[0]
-image
+# Launch the app
+iface.launch()