Update app.py

app.py

@@ -4,66 +4,60 @@ import torch
 from diffusers import FluxPipeline, FluxTransformer2DModel, FlowMatchEulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 from PIL import Image
+import requests
+from translatepy import Translator
 import numpy as np
 import random
 import os
-
 hf_token = os.environ.get('HF_TOKEN')
+from io import BytesIO
+
+translator = Translator()
 
 # Constants
 model = "black-forest-labs/FLUX.1-dev"
+
+
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-@spaces.GPU()
-def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    
-    # Initialize model inside the GPU-enabled function
-    try:
-        transformer = FluxTransformer2DModel.from_single_file(
-            "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v27.safetensors",
-            torch_dtype=torch.bfloat16,
-            token=hf_token
-        )
-    except KeyError as e:
-        print(f"Error loading chroma-unlocked-v27.safetensors: {e}. Falling back to pretrained model.")
-        transformer = FluxTransformer2DModel.from_pretrained(
-            "lodestones/Chroma",
-            subfolder="transformer",
-            torch_dtype=torch.bfloat16,
-            token=hf_token
-        )
-    
+# Ensure model and scheduler are initialized in GPU-enabled function
+if torch.cuda.is_available():
+    transformer = FluxTransformer2DModel.from_single_file(
+        "https://huggingface.co/ekt1701/Test_case/blob/main/rayflux_v10.safetensors",
+        torch_dtype=torch.bfloat16
+    )
     pipe = FluxPipeline.from_pretrained(
-        model,
+        model, 
         transformer=transformer,
-        torch_dtype=torch.bfloat16,
-        token=hf_token
-    )
+        torch_dtype=torch.bfloat16, token=hf_token)
     pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
         pipe.scheduler.config, use_beta_sigmas=True
     )
-    pipe.to(device)
-
-    # Generate images
+    pipe.to("cuda")
+    
+    
+@spaces.GPU()
+def infer(prompt, width, height, num_inference_steps, guidance_scale, nums, seed=42, randomize_seed=True, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    images = pipe(
-        prompt=prompt,
-        width=width,
-        height=height,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=nums,
-        generator=generator
+    generator = torch.Generator().manual_seed(seed)
+    image = pipe(
+            prompt = prompt, 
+            width = width,
+            height = height,
+            num_inference_steps = num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=nums,
+            generator = generator
     ).images
 
-    return images, seed
+    
+    return image, seed
 
-css = """
+
+css="""
 #col-container {
     margin: 0 auto;
     max-width: 1024px;
@@ -71,10 +65,12 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
+    
     with gr.Column(elem_id="col-container"):
-        gr.HTML("<h1><center>Model Testing</center></h1><p><center>Chroma</center></p>")
+        gr.HTML("<h1><center>Image Model Testing</center></h1><p><center>RayFlux V1 Model.</center></p>")
         
         with gr.Row():
+            
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
@@ -82,12 +78,15 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
+            
             run_button = gr.Button("Run", scale=0)
         
-        result = gr.Gallery(label="Gallery", format="png", columns=1, preview=True, height=400)
+        result = gr.Gallery(label="Gallery", format="png", columns = 1, preview=True, height=400)
         
         with gr.Accordion("Advanced Settings", open=False):
+                       
             with gr.Row():
+                
                 width = gr.Slider(
                     label="Width",
                     minimum=256,
@@ -95,6 +94,7 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024,
                 )
+                
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
@@ -104,6 +104,7 @@ with gr.Blocks(css=css) as demo:
                 )
             
             with gr.Row():
+                  
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
@@ -111,6 +112,7 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=30,
                 )
+
                 guidance_scale = gr.Slider(
                     label="Guidance Scale",
                     minimum=0,
@@ -119,7 +121,9 @@ with gr.Blocks(css=css) as demo:
                     value=3.5,
                 )
 
+
             with gr.Row():
+            
                 nums = gr.Slider(
                     label="Number of Images",
                     minimum=1,
@@ -128,6 +132,7 @@ with gr.Blocks(css=css) as demo:
                     value=1,
                     scale=1,
                 )
+       
                 seed = gr.Slider(
                     label="Seed",
                     minimum=0,
@@ -135,13 +140,14 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=-1,
                 )
+            
                 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
-        outputs=[result, seed]
+        fn = infer,
+        inputs = [prompt, width, height, num_inference_steps, guidance_scale, nums, seed, randomize_seed],
+        outputs = [result, seed]
     )
-
+    
 demo.launch()