Update app.py

app.py

@@ -3,13 +3,13 @@ import torch
 import numpy as np
 import modin.pandas as pd
 from PIL import Image
-from diffusers import CogView4Pipeline
+from diffusers import FluxPipeline
 from huggingface_hub import hf_hub_download
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 torch.cuda.max_memory_allocated(device=device)
 torch.cuda.empty_cache()
-pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16)
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
 
 # Open it for reduce GPU memory usage
 pipe.enable_model_cpu_offload()
@@ -24,8 +24,8 @@ def genie (Model, Prompt, negative_prompt, height, width, scale, steps, seed):
         guidance_scale=scale,
         num_images_per_prompt=1,
         num_inference_steps=steps,
-        width=768,
-        height=768,).images[0]
+        width=width,
+        height=height, generator=generator).images[0]
 
     return image